from __future__ import annotations

from typing import Any

import gradio as gr

from src.ui.callbacks import AppCallbacks, current_categories
from src.ui.constants import (
    ALL_CATEGORIES_LABEL,
    APP_NAME,
    APP_SUBTITLE,
    DEFAULT_MIN_SIMILARITY,
    DEFAULT_TOP_K,
    EXAMPLE_QUERIES,
    IMAGE_QUERY_EMPTY_PREVIEW,
    MAX_TOP_K,
    NO_FAVORITES_STATUS,
    TEXT_QUERY_EMPTY_PREVIEW,
    gradio_launch_options,
)
from src.ui.formatters import initial_overview


def create_demo(callbacks: AppCallbacks) -> tuple[gr.Blocks, dict[str, Any]]:
    blocks_kwargs, launch_kwargs = gradio_launch_options()

    with gr.Blocks(**blocks_kwargs) as demo:
        current_results_state = gr.State([])
        selected_result_state = gr.State(None)
        favorites_state = gr.State([])
        query_mode_state = gr.State("text")

        with gr.Column(elem_id="app-shell"):
            gr.Markdown(
                f"# {APP_NAME}\n\n{APP_SUBTITLE}",
                elem_id="hero-section",
            )

            with gr.Row(elem_id="main-dashboard", equal_height=False):
                with gr.Column(scale=1, min_width=300, elem_id="left-sidebar"):
                    with gr.Column(elem_id="search-card", elem_classes=["search-card"]):
                        with gr.Tabs(elem_id="query-tabs"):
                            with gr.Tab("Text Search") as text_tab:
                                text_query = gr.Textbox(
                                    label="Image Description",
                                    placeholder="For example: red apple, milk bottle, green vegetables",
                                    lines=2,
                                    elem_id="text-query",
                                )
                                gr.Examples(
                                    examples=EXAMPLE_QUERIES,
                                    inputs=text_query,
                                    label="Example Queries",
                                    elem_id="example-chips",
                                )
                                text_preview = gr.Markdown(
                                    TEXT_QUERY_EMPTY_PREVIEW,
                                    elem_id="text-preview",
                                )

                            with gr.Tab("Image Search") as image_tab:
                                image_query = gr.Image(
                                    label="Upload Query Image",
                                    type="pil",
                                    placeholder="Drag an image here, or click to upload",
                                    sources=["upload", "clipboard"],
                                    height=210,
                                    elem_id="image-query",
                                )
                                image_preview = gr.Markdown(
                                    IMAGE_QUERY_EMPTY_PREVIEW,
                                    elem_id="image-preview",
                                )

                    with gr.Column(elem_id="parameter-card", elem_classes=["parameter-card"]):
                        gr.Markdown("### Search Parameters", elem_id="parameter-title")
                        top_k = gr.Slider(
                            minimum=1,
                            maximum=MAX_TOP_K,
                            value=DEFAULT_TOP_K,
                            step=1,
                            label="Number of Results",
                            elem_id="top-k",
                        )
                        category_filter = gr.Dropdown(
                            choices=current_categories(),
                            value=ALL_CATEGORIES_LABEL,
                            label="Category Filter",
                            allow_custom_value=False,
                            elem_id="category-filter",
                        )
                        min_similarity = gr.Slider(
                            minimum=0,
                            maximum=1,
                            value=DEFAULT_MIN_SIMILARITY,
                            step=0.01,
                            label="Minimum Similarity",
                            elem_id="min-similarity",
                        )
                        with gr.Row(elem_id="action-row"):
                            search_button = gr.Button(
                                "Search",
                                variant="primary",
                                elem_id="search-button",
                                elem_classes=["primary-button"],
                            )
                            clear_button = gr.Button(
                                "Clear Query",
                                elem_id="clear-button",
                                elem_classes=["secondary-button"],
                            )

                with gr.Column(scale=3, min_width=360, elem_id="right-workspace"):
                    with gr.Column(elem_id="result-card", elem_classes=["result-card"]):
                        result_overview = gr.Markdown(
                            initial_overview(),
                            elem_id="result-overview",
                        )
                        result_gallery = gr.Gallery(
                            label="Search Results",
                            columns=4,
                            height=390,
                            object_fit="cover",
                            show_label=True,
                            elem_id="result-gallery",
                            elem_classes=["result-gallery"],
                        )

                    with gr.Column(elem_id="favorite-card", elem_classes=["favorite-card"]):
                        with gr.Row(elem_id="favorite-header"):
                            gr.Markdown("### My Favorites", elem_id="favorite-title")
                            favorite_status = gr.Markdown(NO_FAVORITES_STATUS, elem_id="favorite-status")
                        with gr.Row(elem_id="favorite-actions"):
                            add_favorite_button = gr.Button(
                                "Add to Favorites",
                                elem_id="add-favorite-button",
                                elem_classes=["secondary-button"],
                            )
                            export_button = gr.DownloadButton(
                                label="Export Favorites as ZIP",
                                value=None,
                                elem_id="export-button",
                                elem_classes=["primary-button"],
                            )
                        favorite_gallery = gr.Gallery(
                            label="My Favorites",
                            columns=5,
                            height=180,
                            object_fit="cover",
                            show_label=True,
                            elem_id="favorite-gallery",
                            elem_classes=["favorite-gallery"],
                        )

        text_tab.select(callbacks.select_text_mode, outputs=query_mode_state)
        image_tab.select(callbacks.select_image_mode, outputs=query_mode_state)
        text_query.change(callbacks.preview_text_query, inputs=text_query, outputs=text_preview)
        image_query.change(callbacks.preview_image_query, inputs=image_query, outputs=image_preview)

        search_button.click(
            callbacks.run_search,
            inputs=[query_mode_state, text_query, image_query, top_k, category_filter, min_similarity],
            outputs=[result_overview, result_gallery, selected_result_state, current_results_state],
        )
        clear_button.click(
            callbacks.clear_search,
            inputs=[top_k, category_filter, min_similarity],
            outputs=[
                text_query,
                image_query,
                text_preview,
                image_preview,
                result_overview,
                result_gallery,
                selected_result_state,
                current_results_state,
            ],
        )
        result_gallery.select(
            callbacks.select_result_from_gallery,
            inputs=current_results_state,
            outputs=selected_result_state,
            show_progress="hidden",
        )
        add_favorite_button.click(
            callbacks.add_selected_to_favorites,
            inputs=[selected_result_state, current_results_state, favorites_state],
            outputs=[favorites_state, favorite_gallery, favorite_status, export_button],
        )

    return demo, launch_kwargs