import time
from pathlib import Path

import gradio as gr

from qa import ask_about_image
from search import search, filenames
from ui.theme import theme, CSS, TOGGLE_DARK


# ============================================================
# Config
# ============================================================

IMAGE_DIR = Path("data/Flicker8k_Dataset")
EXAMPLES = [
    "a dog running on the beach",
    "children playing soccer",
    "a man climbing a rock",
    "a person riding a bike in the snow",
]

SCROLL_TO_ASK = """
() => document.getElementById('ask-panel')
        ?.scrollIntoView({ behavior: 'smooth', block: 'center' })
"""

ANSWER_PLACEHOLDER = "Ask a question about this image."

# ============================================================
# Handlers
# ============================================================

def run_search(query: str, k: int):
    if not query.strip():
        return (gr.update(value=[], visible=False), "", gr.update(visible=True),
                [], gr.update(visible=False))

    start = time.perf_counter()
    results = search(query, k=int(k))
    elapsed = time.perf_counter() - start

    paths, items = [], []
    for rank, (name, score) in enumerate(results, start=1):
        path = str(IMAGE_DIR / name)
        paths.append(path)
        items.append((path, f"#{rank} · {score:.3f}"))

    status_text = f"{len(results)} of {len(filenames):,} · {elapsed:.2f} s"
    return (gr.update(value=items, visible=True), status_text, gr.update(visible=False),
            paths, gr.update(visible=False))


def on_select(evt: gr.SelectData, paths: list[str]):
    path = paths[evt.index]
    info = f"Result #{evt.index + 1} · `{Path(path).name}`"
    return gr.update(visible=True), path, path, info, "", ANSWER_PLACEHOLDER

def on_ask(image_path: str | None, question: str):
    if not image_path:
        return "Click an image in the results first."
    return ask_about_image(image_path, question)

# ============================================================
# Layout
# ============================================================

with gr.Blocks(title="Image Search") as demo:
    result_paths = gr.State([])
    selected_path = gr.State(None)

    # --- Header ---
    with gr.Row(elem_id="header"):
        gr.Markdown(
            "# Image Search\nDescribe a scene. We rank the photos by similarity.",
            elem_id="title",
        )
        theme_btn = gr.Button("Light / Dark", size="sm", variant="secondary", elem_id="theme-btn")

    # --- Search ---
    with gr.Column(elem_classes="card"):
        with gr.Row(elem_id="search-row"):
            query = gr.Textbox(label="Query", placeholder="a dog running on the beach", scale=5)
            k = gr.Slider(1, 20, value=8, step=1, label="Top-k", scale=2)
            button = gr.Button("Search", variant="primary", scale=1, min_width=120, elem_id="search-btn")
        with gr.Row(elem_id="chips"):
            gr.Markdown("Try:", elem_id="try-label")
            chips = [gr.Button(ex, size="sm", variant="secondary", elem_classes="chip") for ex in EXAMPLES]

    # --- Results ---
    with gr.Row(elem_id="results-head"):
        gr.Markdown("## Results")
        status = gr.Markdown("", elem_id="status")

    empty = gr.Markdown("Your search results will appear here", elem_id="empty-state")

    gallery = gr.Gallery(
        show_label=False,
        container=False,
        columns=4,
        height="auto",
        object_fit="cover",
        allow_preview=False,
        elem_id="gallery",
        visible=False,
    )

    # --- Ask about this image ---
    with gr.Column(elem_classes="card", elem_id="ask-panel", visible=False) as ask_panel:
        with gr.Row(elem_id="ask-row"):
            selected_img = gr.Image(
                show_label=False, interactive=False, height=220,
                scale=1, min_width=220, elem_id="ask-image",
            )
            with gr.Column(scale=3):
                gr.Markdown("### Ask about this image")
                selected_info = gr.Markdown("", elem_id="ask-info")
                with gr.Row(elem_id="ask-input-row"):
                    question = gr.Textbox(
                        show_label=False, placeholder="e.g. What color is the dog?",
                        scale=5,
                    )
                    ask_btn = gr.Button("Ask", variant="primary", scale=1, min_width=100)
                answer = gr.Markdown(ANSWER_PLACEHOLDER, elem_id="ask-answer")


    # ========================================================
    # Events
    # ========================================================

    search_outputs = [gallery, status, empty, result_paths, ask_panel]

    theme_btn.click(None, js=TOGGLE_DARK)
    button.click(run_search, inputs=[query, k], outputs=search_outputs)
    query.submit(run_search, inputs=[query, k], outputs=search_outputs)
    for chip, ex in zip(chips, EXAMPLES):
        chip.click(lambda q=ex: q, outputs=query).then(
            run_search, inputs=[query, k], outputs=search_outputs
        )

    gallery.select(
        on_select, inputs=[result_paths],
        outputs=[ask_panel, selected_img, selected_path, selected_info, question, answer],
    ).then(None, js=SCROLL_TO_ASK)

    ask_btn.click(on_ask, inputs=[selected_path, question], outputs=answer, concurrency_limit=2)
    question.submit(on_ask, inputs=[selected_path, question], outputs=answer, concurrency_limit=2)

# ============================================================
# Run
# ============================================================

if __name__ == "__main__":
    demo.launch(theme=theme, css=CSS, server_name="0.0.0.0")