import gradio as gr


from src.generate_queries_alternative import main_generate_queries
import time
import pandas as pd

import spacy

# Try to load the model, and download it if missing
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    from spacy.cli import download
    download("en_core_web_sm")
    nlp = spacy.load("en_core_web_sm")

from src.visual_qa import main_text_retrieve_images

def handle_structured_query(question, sort_by=""):
    if not question:
        return "Please ask something 🙂", pd.DataFrame(), []

    try:
        start = time.time()
        result_query, sparql_query = main_generate_queries(question)
        elapsed = round(time.time() - start, 2)
    except Exception as e:
        return f"⚠️ Query failed: {e}", pd.DataFrame(), []

    if isinstance(result_query, str):
        return result_query, pd.DataFrame(), []

    if not result_query:
        return f"No results for '{question}'. Try rephrasing. (⏱ {elapsed}s)", pd.DataFrame(), []

    df = pd.DataFrame(result_query)
    if sort_by and sort_by in df.columns:
        df = df.sort_values(by=sort_by)


    if "image_url" in df.columns:
        columns_of_interest = ["image_url", "year","fashion_collectionLabel", "reference_URL"]
        df = df[columns_of_interest]
        # Create a gallery: each item is (image_url, metadata string)
        gallery_items = []
        for _, row in df.iterrows():
            image_url = row.get("image_url")
            if not image_url:
                continue
            # Caption from other fields
            caption = " | ".join(f"{k}: {v}" for k, v in row.items() if k != "image_url" and pd.notnull(v))
            gallery_items.append((image_url, caption))
        return f"Query returned {len(gallery_items)} image(s) in {elapsed} seconds.", pd.DataFrame(), gallery_items

    return f"Query returned a table with {len(df)} row(s) in {elapsed} seconds.", df, []




def handle_image_query(text):
    if not text:
        return []

    try:
        records = main_text_retrieve_images(text)
    except Exception as e:
        return [("https://via.placeholder.com/300x200?text=Error", f"Error: {e}")]

    gallery_items = []
    for item in records:
        image_url = item.get("image_url")
        if not image_url:
            continue
        # Build a simple caption from the remaining fields
        caption = " | ".join(f"{k}: {v}" for k, v in item.items() if k != "image_url")
        gallery_items.append((image_url, caption))

    return gallery_items

# --- UI --- #
with gr.Blocks() as demo:
    gr.Markdown("# 🧵 FashionDB Interface")


    with gr.Tab("Structured Query"):
        gr.Markdown("Ask FashionDB anything and view results with images + metadata.")

        with gr.Row():
            query_input = gr.Textbox(label="Your question")
            sort_input = gr.Textbox(label="Sort by (optional column name)", placeholder="e.g. start_year")

        query_submit = gr.Button("Submit")

        query_text_output = gr.Textbox(label="Message", interactive=False)
        query_table_output = gr.Dataframe(label="Tabular Result", interactive=False)
        query_gallery_output = gr.Gallery(label="Image Gallery")
        query_submit.click(
            fn=handle_structured_query,
            inputs=[query_input, sort_input],
            outputs=[
                query_text_output,
                query_table_output,
                query_gallery_output
            ]
        )

    with gr.Tab("Image Retrieval"):
        gr.Markdown("Search for similar fashion show images based on a text description.")
        image_text = gr.Textbox(label="Describe the kind of images you're looking for")
        image_submit = gr.Button("Find Images")
        image_gallery = gr.Gallery(label="Retrieved Images")

        image_submit.click(handle_image_query, inputs=image_text, outputs=image_gallery)

demo.launch()