import gradio as gr
import os
import json
import uuid
from datetime import datetime
import shutil
from huggingface_hub import HfApi, create_repo, upload_file, upload_folder
os.makedirs("uploaded_images", exist_ok=True)
os.makedirs("submissions", exist_ok=True)
HF_TOKEN = os.environ.get("Crowdsourcing") 
DATASET_NAME = "1-800-LLMs/se-culture-dataset-results" 
DATASET_CREATED = False
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", "Khulna", "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"]}
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"]
def setup_hf_dataset():
    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")
        except Exception as e: print(f"Error setting up dataset: {e}")
    elif not HF_TOKEN:
        print("Warning: HF_TOKEN not set. Data will be stored locally only.")
def update_state_dropdown(country):
    if country in states_by_country:
        return gr.Dropdown(choices=states_by_country[country], label=f"State/Province in {country}:", interactive=True)
    return gr.Dropdown(choices=[], label="State/Province:", interactive=True)
def update_other_language_visibility(selected_language):
    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,code_mixed_caption,domain,email,other_language=None):
    submission_id = str(uuid.uuid4())
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    image_path = None
    if input_img is not None:
        image_filename = f"{timestamp}.jpg"
        image_path = os.path.join("uploaded_images", image_filename)
        if isinstance(input_img, str): shutil.copy(input_img, image_path)
        else: input_img.save(image_path)
    submission_data = {"id": submission_id, "timestamp": timestamp, "image_filename": os.path.basename(image_path) if image_path else None, "language": language, "country": country, "state": state, "city": city, "se_asia_relevance": se_asia_relevance, "cultural_knowledge_source": culture_knowledge, "native_caption": native_caption, "english_caption": english_caption,"code_mixed_caption": code_mixed_caption,"domain":domain, "email": email}
    if other_language is not None:
        submission_data["other language"] = other_language
    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 and DATASET_CREATED:
        try:
            api = HfApi()
            api.upload_file(path_or_fileobj=json_path, path_in_repo=f"submissions/{json_filename}", repo_id=DATASET_NAME, repo_type="dataset", token=HF_TOKEN)
            if image_path and os.path.exists(image_path):
                api.upload_file(path_or_fileobj=image_path, path_in_repo=f"images/{os.path.basename(image_path)}", repo_id=DATASET_NAME, repo_type="dataset", token=HF_TOKEN)
            print(f"Submission {submission_id} uploaded to Hugging Face Dataset")
        except Exception as e: print(f"Error uploading to dataset: {e}")
    location_info = f"Location: {city}, {state}, {country}" if state else f"Location: {city}, {country}"
    response = [input_img, f"Language: {language}", f"Selected location: {location_info}", f"SE Asia relevance: {se_asia_relevance}", f"Cultural knowledge source: {culture_knowledge}", f"Native caption: {native_caption}", f"English caption: {english_caption}" , f"Code mixed caption: {code_mixed_caption}", f"Domain: {domain}"]
    if other_language is not None:
        response.append(f"Other language: {other_language}")
    return tuple(response)
def clear_inputs():
    return None, "OTHER", "", "None", None, "", "", None, None, "", "", "", ""
setup_hf_dataset()
with gr.Blocks(theme='1024m/1024m-1') as gradio_app:
    gr.Markdown("# Multilingual Image Captions")
    gr.Markdown("Please make sure to check the [annotation guidelines](https://www.google.com/) and the [discord channel](https://www.google.com/) before proceeding.")
    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(label="Upload an image", sources=['upload', 'webcam'], type="pil")
            language = gr.Dropdown(choices=south_asian_languages, label="Language:", info="Select the native language relevant to the image", interactive=True,value=south_asian_languages[-1])
            other_language = gr.Textbox(label="Other Language:", info="Name of the language if not listed above", visible=True)
            country_dropdown = gr.Dropdown(choices=["None","India", "Pakistan", "Bangladesh", "Afghanistan", "Bhutan", "Nepal", "Sri Lanka"], label="Country where the image was taken:", interactive=True)
            state_dropdown = gr.Dropdown(choices=[], label="State/Province:", interactive=True)
            city_textbox = gr.Textbox(label="City where the image was taken:", placeholder="Enter city name")
            email_input = gr.Textbox(label="Your Email:", placeholder="Enter your email address", info="Used as unique contributor ID")
        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?")
            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.)")
            native_caption = gr.Textbox(label="Caption in Native Language:", placeholder="Enter caption in the native language of the culture depicted", info="in native script ONLY")
            english_caption = gr.Textbox(label="English Caption:", placeholder="Enter caption in English", info="in english script ONLY")
            code_mixed_caption = gr.Textbox(label="Code-Mixed Caption:", placeholder="Enter caption in code-mixed", info="in english script ONLY")
            domain = gr.Textbox(label="Domain:",placeholder="Description",info="1-2 word any 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")
            output_text = gr.Text(label="Text Response")
            output_location = gr.Text(label="Location Information")
        with gr.Column(scale=1):
            output_relevance = gr.Text(label="South Asia Cultural Relevance")
            output_knowledge = gr.Text(label="Cultural Knowledge Source")
            output_native = gr.Text(label="Native Language Caption")
            output_english = gr.Text(label="English Caption")
    language.change(update_other_language_visibility, inputs=language, outputs=other_language)
    country_dropdown.change(fn=update_state_dropdown, inputs=country_dropdown, outputs=state_dropdown)
    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,code_mixed_caption,domain, email_input,other_language], outputs=[output_img, output_text, output_location, output_relevance, output_knowledge, output_native, output_english])
    clear_btn.click(fn=clear_inputs, inputs=[], outputs=[input_img, language,other_language, country_dropdown, state_dropdown, city_textbox, se_asia_relevance, culture_knowledge, native_caption, english_caption,code_mixed_caption,domain, email_input])
if __name__ == "__main__":
    gradio_app.launch()