import gradio as gr
import os
import json
import uuid
from datetime import datetime
from huggingface_hub import HfApi, create_repo

# Setup local storage
os.makedirs("uploaded_images", exist_ok=True)
os.makedirs("submissions", exist_ok=True)

# --- Hugging Face Configuration ---
HF_TOKEN = os.environ.get("Crowdsourcing")
DATASET_NAME = "1-800-LLMs/se-culture-dataset-results"
DATASET_CREATED = False

# --- Data for Dropdowns ---
states_by_country = {
    "India": ["Andhra Pradesh", "Arunachal Pradesh", "Assam", "Bihar", "Chhattisgarh", "Goa", "Gujarat", "Haryana","Himachal Pradesh", "Jharkhand", "Karnataka", "Kerala", "Madhya Pradesh", "Maharashtra", "Manipur","Meghalaya", "Mizoram", "Nagaland", "Odisha", "Punjab", "Rajasthan", "Sikkim", "Tamil Nadu","Telangana", "Tripura", "Uttar Pradesh", "Uttarakhand", "West Bengal", "Andaman and Nicobar Islands","Chandigarh", "Dadra and Nagar Haveli and Daman and Diu", "Delhi", "Jammu and Kashmir", "Ladakh","Lakshadweep", "Puducherry"],
    "Pakistan": ["Balochistan", "Khyber Pakhtunkhwa", "Punjab", "Sindh", "Islamabad Capital Territory", "Other"],
    "Bangladesh": ["Barisal", "Chittagong", "Dhaka", "Khulnā", "Mymensingh", "Rajshahi", "Rangpur", "Sylhet"],
    "Afghanistan": ["Badakhshan", "Badghis", "Baghlan", "Balkh", "Bamyan", "Daykundi", "Farah", "Faryab", "Ghazni","Ghor", "Helmand", "Herat", "Jowzjan", "Kabul", "Kandahar", "Kapisa", "Khost", "Kunar", "Kunduz","Laghman", "Logar", "Nangarhar", "Nimruz", "Nuristan", "Paktia", "Paktika", "Panjshir", "Parwan","Samangan", "Sar-e Pol", "Takhar", "Uruzgan", "Wardak", "Zabul"],
    "Bhutan": ["Bumthang", "Chukha", "Dagana", "Gasa", "Haa", "Lhuentse", "Mongar", "Paro", "Pemagatshel", "Punakha","Samdrup Jongkhar", "Samtse", "Sarpang", "Thimphu", "Trashigang", "Trashiyangtse", "Trongsa", "Tsirang","Wangdue Phodrang", "Zhemgang"],
    "Nepal": ["Bagmati", "Gandaki", "Karnali", "Koshi", "Lumbini", "Madhesh", "Sudurpashchim"],
    "Sri Lanka": ["Central", "Eastern", "North Central", "Northern", "North Western", "Sabaragamuwa", "Southern","Uva", "Western"]
}
countries = ["India", "Pakistan", "Bangladesh", "Afghanistan", "Bhutan", "Nepal", "Sri Lanka","OTHER"]

south_asian_languages = [
    "Assamese", "Bengali", "Bhojpuri", "Bodo", "Dari", "Dzongkha", "Dogri", "Gujarati", "Hindi", "Kannada",
    "Kashmiri", "Konkani", "Maithili", "Malayalam", "Marathi", "Meitei", "Nepali", "Odia", "Pashto", "Punjabi",
    "Sanskrit", "Santali", "Sindhi", "Sinhala", "Tamil", "Telugu", "Tibetan", "Tulu", "Urdu", "OTHER"
]

# --- Helper Functions ---

def setup_hf_dataset():
    """Creates the Hugging Face dataset repository if it doesn't exist."""
    global DATASET_CREATED
    if not DATASET_CREATED and HF_TOKEN:
        try:
            api = HfApi()
            create_repo(DATASET_NAME, repo_type="dataset", token=HF_TOKEN, exist_ok=True)
            DATASET_CREATED = True
            print(f"Dataset {DATASET_NAME} is ready on Hugging Face Hub.")
        except Exception as e:
            print(f"Error setting up Hugging Face dataset: {e}")
    elif not HF_TOKEN:
        print("Warning: HF_TOKEN not set. Submissions will be stored locally only.")

def update_country_dependents(country):
    """
    Updates the state dropdown, other state textbox, and other country textbox
    based on the selected country.
    """
    # Logic for the state dropdown
    if country in states_by_country:
        state_update = gr.update(
            choices=states_by_country[country],
            label=f"State/Province in {country}:",
            interactive=True,
            value=None, # Reset selection
            visible=True
        )
        other_state_update = gr.update(visible=False, value="") # Hide other state textbox
    else:
        # Hide state dropdown for "None" or "OTHER"
        state_update = gr.update(
            choices=[],
            interactive=False,
            value=None,
            visible=False
        )
        # Show 'Other State' textbox ONLY if country is 'OTHER'
        other_state_update = gr.update(visible=(country == "OTHER"))

    # Logic for the 'Other Country' textbox visibility
    other_country_update = gr.update(visible=(country == "OTHER"))

    # Return updates for all three components
    return state_update, other_country_update, other_state_update


def update_other_language_visibility(selected_language):
    """Shows the 'Other Language' textbox only when 'OTHER' is selected."""
    return gr.update(visible=(selected_language == "OTHER"))

def process_submission(input_img, language, country, state, city, se_asia_relevance, culture_knowledge,
                       native_caption, english_caption, transliterated_caption, domain, email, other_language, other_country, other_state):
    """Validates, saves, and uploads a user submission."""
    warnings = {
        "img": (not input_img, "<span style='color:red'>⚠️ Please upload an image.</span>"),
        "lang": (not language or (language == "OTHER" and not other_language), "<span style='color:red'>⚠️ Please select or specify a language.</span>"),
        "country": (not country or country == "None" or (country == "OTHER" and not other_country), "<span style='color:red'>⚠️ Please select or specify a country.</span>"),
        "email": (not email, "<span style='color:red'>⚠️ Please provide your email.</span>"),
        "relevance": (not se_asia_relevance, "<span style='color:red'>⚠️ Please select the cultural relevance.</span>"),
        "knowledge": (not culture_knowledge, "<span style='color:red'>⚠️ Please select your knowledge source.</span>"),
        "english": (not english_caption, "<span style='color:red'>⚠️ Please enter an English caption.</span>"),
        "transliterated": (not transliterated_caption, "<span style='color:red'>⚠️ Please enter a transliterated caption.</span>"),
    }
    if any(v[0] for v in warnings.values()):
        return (
            gr.update(visible=True, value=warnings["img"][1] if warnings["img"][0] else ""),
            gr.update(visible=True, value=warnings["lang"][1] if warnings["lang"][0] else ""),
            gr.update(visible=True, value=warnings["country"][1] if warnings["country"][0] else ""),
            gr.update(visible=False),
            gr.update(visible=True, value=warnings["email"][1] if warnings["email"][0] else ""),
            gr.update(visible=True, value=warnings["relevance"][1] if warnings["relevance"][0] else ""),
            gr.update(visible=True, value=warnings["knowledge"][1] if warnings["knowledge"][0] else ""),
            gr.update(visible=False),
            gr.update(visible=True, value=warnings["english"][1] if warnings["english"][0] else ""),
            gr.update(visible=True, value=warnings["transliterated"][1] if warnings["transliterated"][0] else ""),
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
            gr.update(visible=False), gr.update(visible=False)
        )
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    image_filename = f"{timestamp}.jpg"
    image_path = os.path.join("uploaded_images", image_filename)
    input_img.save(image_path)
    
    final_language = other_language if language == "OTHER" else language
    final_country = other_country if country == "OTHER" else country
    final_state = other_state if country == "OTHER" else state
    
    submission_data = {
        "id": str(uuid.uuid4()), "timestamp": timestamp, "image_filename": image_filename,
        "language": final_language, "country": final_country, "state": final_state, "city": city,
        "se_asia_relevance": se_asia_relevance, "cultural_knowledge_source": culture_knowledge,
        "native_caption": native_caption, "english_caption": english_caption,
        "transliterated_caption": transliterated_caption, "domain": domain, "email": email,
    }
    json_filename = f"{timestamp}.json"
    json_path = os.path.join("submissions", json_filename)
    with open(json_path, "w") as f:
        json.dump(submission_data, f, indent=2)
        
    if HF_TOKEN:
        try:
            api = HfApi()
            api.upload_file(path_or_fileobj=json_path, path_in_repo=f"data/{json_filename}", repo_id=DATASET_NAME, repo_type="dataset", token=HF_TOKEN)
            api.upload_file(path_or_fileobj=image_path, path_in_repo=f"images/{image_filename}", repo_id=DATASET_NAME, repo_type="dataset", token=HF_TOKEN)
            print(f"Successfully uploaded {json_filename} and {image_filename} to HF.")
        except Exception as e:
            print(f"Error uploading to Hugging Face Hub: {e}")
            
    location_info = f"Location: {city}, {final_state}, {final_country}" if city and final_state else f"Location: {final_country}"
    
    return (
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value=input_img, visible=True), gr.update(value="✅ Submission successful!", visible=True),
        gr.update(value=location_info, visible=True), gr.update(value=se_asia_relevance, visible=True),
        gr.update(value=culture_knowledge, visible=True), gr.update(value=native_caption, visible=True),
        gr.update(value=english_caption, visible=True), gr.update(value=domain, visible=True)
    )

def clear_all_fields():
    """Resets all input and output components in the UI."""
    return (
        None, "OTHER", "", "None", "", None, "", "", # Add "" for other_state_textbox
        None, None, "", "", "", "", "",
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(value="", visible=False), gr.update(value="", visible=False),
        gr.update(visible=False, value=None), gr.update(visible=False, value=None),
        gr.update(visible=False, value=None), gr.update(visible=False, value=None),
        gr.update(visible=False, value=None), gr.update(visible=False, value=None),
        gr.update(visible=False, value=None), gr.update(visible=False, value=None)
    )

# --- Main Application ---

setup_hf_dataset()

with gr.Blocks(theme='1024m/1024m-1') as gradio_app:
    gr.Markdown("# Multilingual Image Captions")
    gr.Markdown("Please check the [annotation guidelines](https://docs.google.com/document/d/1GiPRvIP44EMl1lSjREegfYORqXZzdVufBnMHTogNGfw/edit?usp=sharing) and the [discord channel](https://discord.com/channels/987824841656791130/991440413745500210) before proceeding.")

    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(label="Upload an image", sources=['upload', 'webcam'], type="pil")
            img_warning = gr.Markdown(visible=False)
            language = gr.Dropdown(choices=south_asian_languages, label="Language:", interactive=True, value=south_asian_languages[-1])
            lang_warning = gr.Markdown(visible=False)
            other_language = gr.Textbox(label="Other Language:", info="If not listed above", visible=True)
            countries = ["India", "Pakistan", "Bangladesh", "Afghanistan", "Bhutan", "Nepal", "Sri Lanka", "OTHER"]
            country_dropdown = gr.Dropdown(choices=countries, label="Country where the image was taken:", interactive=True, value=countries[-1])
            country_warning = gr.Markdown(visible=False)
            other_country = gr.Textbox(label="Other Country:", info="If not listed above", visible=True)
            state_dropdown = gr.Dropdown(choices=[], label="State/Province (Optional but preferred):", interactive=False, visible=True)
            other_state_textbox = gr.Textbox(label="Other State/Province:", info="If country is 'OTHER'", visible=False)
            city_textbox = gr.Textbox(label="City (Optional but preferred):", placeholder="Enter city name")
            city_warning = gr.Markdown(visible=False)
            email_input = gr.Textbox(label="Your Email:", placeholder="Enter your email address", info="Used as unique contributor ID")
            email_warning = gr.Markdown(visible=False)
        with gr.Column(scale=1):
            se_asia_relevance = gr.Radio(choices=["Yes. Unique to South Asia", "Yes, people will likely think of South Asia when seeing the picture, but it may have low degree of similarity to other cultures.", "Maybe, this culture did not originate from South Asia, but it's quite dominant in South Asia", "Not really. It has some affiliation to South Asia, but actually does not represent South Asia or has stronger affiliation to cultures outside South Asia", "No. Totally unrelated to South Asia"], label="Is the image culturally relevant in South Asia?")
            relevance_warning = gr.Markdown(visible=False)
            culture_knowledge = gr.Radio(choices=["I'm from this country/culture", "I checked online resources (e.g., Wikipedia, articles, blogs)"], label="How do you know about this culture?", info="Please do not consult LLMs (e.g., GPT-4o, Claude, Command-R, etc.)")
            knowledge_warning = gr.Markdown(visible=False)
            native_caption = gr.Textbox(label="Caption in Native Language (Optional but preferred):", placeholder="Enter caption in the native language (script only)")
            native_warning = gr.Markdown(visible=False)
            english_caption = gr.Textbox(label="English Caption:", placeholder="Enter caption in English (script only)")
            english_warning = gr.Markdown(visible=False)
            transliterated_caption = gr.Textbox(label="Transliterated Caption:", placeholder="Enter caption in transliterated (English script only)")
            transliterated_warning = gr.Markdown(visible=False)
            domain = gr.Textbox(label="Domain (Optional but preferred):", placeholder="1-2 word description")

    with gr.Row():
        clear_btn = gr.Button("Clear")
        submit_btn = gr.Button("Submit")

    with gr.Row():
        with gr.Column(scale=1):
            output_img = gr.Image(label="Submitted Image", visible=False)
            output_text = gr.Text(label="Submission Status", visible=False)
            output_location = gr.Text(label="Location Information", visible=False)
        with gr.Column(scale=1):
            output_relevance = gr.Text(label="South Asia Cultural Relevance", visible=False)
            output_knowledge = gr.Text(label="Cultural Knowledge Source", visible=False)
            output_native = gr.Text(label="Native Language Caption", visible=False)
            output_english = gr.Text(label="English Caption", visible=False)
            output_domain = gr.Text(label="Domain", visible=False)
            
    # Event bindings
    language.change(update_other_language_visibility, inputs=language, outputs=other_language)
    
    country_dropdown.change(
        fn=update_country_dependents,
        inputs=country_dropdown,
        outputs=[state_dropdown, other_country, other_state_textbox]
    )

    submit_btn.click(
        fn=process_submission,
        inputs=[
            input_img, language, country_dropdown, state_dropdown, city_textbox,
            se_asia_relevance, culture_knowledge, native_caption, english_caption,
            transliterated_caption, domain, email_input, other_language, other_country,
            other_state_textbox
        ],
        outputs=[
            img_warning, lang_warning, country_warning, city_warning, email_warning,
            relevance_warning, knowledge_warning, native_warning, english_warning, transliterated_warning,
            output_img, output_text, output_location,
            output_relevance, output_knowledge, output_native, output_english, output_domain
        ]
    )

    components_to_clear = [
        input_img, language, other_language, country_dropdown, other_country, state_dropdown,
        other_state_textbox, city_textbox, se_asia_relevance, culture_knowledge, native_caption,
        english_caption, transliterated_caption, domain, email_input,
        img_warning, lang_warning, country_warning, city_warning, email_warning, relevance_warning,
        knowledge_warning, native_warning, english_warning, transliterated_warning,
        output_img, output_text,
        output_location, output_relevance, output_knowledge, output_native, output_english, output_domain
    ]
    clear_btn.click(fn=clear_all_fields, inputs=None, outputs=components_to_clear)

if __name__ == "__main__":
    gradio_app.launch(debug=True)