import gradio as gr
import os
import json
import uuid
from datetime import datetime
import shutil

# Create directories for data storage
os.makedirs("uploaded_images", exist_ok=True)
os.makedirs("submissions", exist_ok=True)

def process_submission(input_img, text_answer, multiple_choice, city, country, se_asia_relevance, culture_knowledge, native_caption, english_caption):
    # Generate unique ID for this submission
    submission_id = str(uuid.uuid4())
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Save the image if provided
    image_path = None
    if input_img is not None:
        # Create filename with submission ID
        image_filename = f"{timestamp}_{submission_id}.jpg"
        image_path = os.path.join("uploaded_images", image_filename)
        
        # Save the image
        if isinstance(input_img, str):  # If it's a file path
            shutil.copy(input_img, image_path)
        else:  # If it's a PIL Image
            input_img.save(image_path)
    
    # Create a data structure for the submission
    submission_data = {
        "id": submission_id,
        "timestamp": timestamp,
        "image_filename": image_path,
        "cultural_relevance": text_answer,
        "continent": multiple_choice,
        "city": city,
        "country": country,
        "se_asia_relevance": se_asia_relevance,
        "cultural_knowledge_source": culture_knowledge,
        "native_caption": native_caption,
        "english_caption": english_caption
    }
    
    # Save the data as JSON
    json_filename = f"{timestamp}_{submission_id}.json"
    json_path = os.path.join("submissions", json_filename)
    with open(json_path, "w") as f:
        json.dump(submission_data, f, indent=2)
    
    # Return values to display in the interface
    return input_img, f"Your text response: {text_answer}", f"Your selected option: {multiple_choice}", f"Location: {city}, {country}", f"SE Asia relevance: {se_asia_relevance}", f"Cultural knowledge source: {culture_knowledge}", f"Native caption: {native_caption}", f"English caption: {english_caption}"

gradio_app = gr.Interface(
    process_submission,
    inputs=[
        gr.Image(label="Upload an image", sources=['upload', 'webcam'], type="pil"),
        gr.Textbox(label="The image portrays culturally-relevant information in:", placeholder="what culture does this image represent?"),
        gr.Radio(choices=["North America", "South America", "Europe", "Africa", "Asia", "Australia"], 
                 label="Which continent is most represented in this image?"),
        gr.Textbox(label="City where the image was taken:", placeholder="Enter city name"),
        gr.Textbox(label="Country where the image was taken:", placeholder="Enter country name"),
        gr.Radio(
            choices=[
                "Yes. Unique to South East",
                "Yes, people will likely think of South East when seeing the picture, but it may have low degree of similarity to other cultures.",
                "Maybe, this culture did not originate from South East, but it's quite dominant in South East",
                "Not really. It has some affiliation to South East, but actually does not represent South East or has stronger affiliation to cultures outside South East",
                "No. Totally unrelated to South East"
            ],
            label="Is the image culturally relevant in South-East Asia?"
        ),
        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.)"
        ),
        gr.Textbox(label="Caption in Native Language:", placeholder="Enter caption in the native language of the culture depicted"),
        gr.Textbox(label="English Caption:", placeholder="Enter caption in English")
    ],
    outputs=[
        gr.Image(label="Submitted Image"),
        gr.Text(label="Text Response"),
        gr.Text(label="Multiple Choice Response"),
        gr.Text(label="Location Information"),
        gr.Text(label="South-East Asia Cultural Relevance"),
        gr.Text(label="Cultural Knowledge Source"),
        gr.Text(label="Native Language Caption"),
        gr.Text(label="English Caption")
    ],
    title="SE Image Data Collection",
    description="Upload an image and answer questions about its cultural significance."
)

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