"""Hugging Face ZeroGPU entrypoint for SatQuery AI."""

from __future__ import annotations

import gradio as gr
import spaces

from satquery_space_bridge import (
    api_contract,
    preload_models,
    run_bitemporal,
    run_optical_sar,
    run_single_image,
    space_health,
)


# ZeroGPU provides CUDA emulation during module initialization. Keep model
# construction here and inference inside the decorated request handlers below.
preload_models()


@spaces.GPU(duration=90)
def single_image_api(
    primary_image: str | None,
    query: str,
    primary_modality: str,
    primary_image_modality: str,
    use_cache: bool,
    force_rerun: bool,
):
    return run_single_image(
        primary_image,
        query,
        primary_modality,
        primary_image_modality,
        use_cache,
        force_rerun,
    )


@spaces.GPU(duration=75)
def optical_sar_api(
    optical_image: str | None,
    sar_image: str | None,
    query: str,
    use_cache: bool,
    force_rerun: bool,
):
    return run_optical_sar(optical_image, sar_image, query, use_cache, force_rerun)


@spaces.GPU(duration=90)
def bitemporal_api(
    earlier_image: str | None,
    later_image: str | None,
    query: str,
    primary_date: str,
    secondary_date: str,
    primary_modality: str,
    use_cache: bool,
    force_rerun: bool,
):
    return run_bitemporal(
        earlier_image,
        later_image,
        query,
        primary_date,
        secondary_date,
        primary_modality,
        use_cache,
        force_rerun,
    )


def _outputs():
    return [
        gr.JSON(label="Result envelope"),
        gr.JSON(label="Artifact manifest"),
        gr.File(label="Evidence files", file_count="multiple"),
    ]


with gr.Blocks(title="SatQuery AI ZeroGPU Bridge") as demo:
    gr.Markdown(
        "# SatQuery AI API bridge\n"
        "This minimal Space exposes the existing SatQuery workflows as named Gradio APIs. "
        "Use **View API** for the generated client contract."
    )

    with gr.Tab("Single Image"):
        single_file = gr.File(label="Image", type="filepath")
        single_query = gr.Textbox(label="Query", value="Describe this satellite image.")
        single_modality = gr.Dropdown(
            ["optical", "multispectral", "sar", "unknown"], value="optical", label="Modality"
        )
        single_image_modality = gr.Dropdown(
            [
                "auto",
                "optical_rgb",
                "optical_grayscale",
                "panchromatic",
                "sar_preview",
                "sar_vv",
                "sar_vh",
                "sar_vv_vh",
                "multispectral",
                "unknown",
            ],
            value="auto",
            label="Image modality",
        )
        single_cache = gr.Checkbox(value=True, label="Use cache")
        single_rerun = gr.Checkbox(value=False, label="Force rerun")
        single_submit = gr.Button("Run single-image workflow", variant="primary")
        single_submit.click(
            single_image_api,
            [single_file, single_query, single_modality, single_image_modality, single_cache, single_rerun],
            _outputs(),
            api_name="single_image",
        )

    with gr.Tab("Optical + SAR"):
        optical_file = gr.File(label="Optical image", type="filepath")
        sar_file = gr.File(label="SAR image", type="filepath")
        cross_query = gr.Textbox(label="Query", value="Compare the optical and SAR observations.")
        cross_cache = gr.Checkbox(value=True, label="Use cache")
        cross_rerun = gr.Checkbox(value=False, label="Force rerun")
        cross_submit = gr.Button("Run optical + SAR workflow", variant="primary")
        cross_submit.click(
            optical_sar_api,
            [optical_file, sar_file, cross_query, cross_cache, cross_rerun],
            _outputs(),
            api_name="optical_sar",
        )

    with gr.Tab("Bi-temporal"):
        earlier_file = gr.File(label="Earlier image", type="filepath")
        later_file = gr.File(label="Later image", type="filepath")
        temporal_query = gr.Textbox(label="Query", value="What changed between these images?")
        earlier_date = gr.Textbox(label="Earlier date", placeholder="YYYY-MM-DD")
        later_date = gr.Textbox(label="Later date", placeholder="YYYY-MM-DD")
        temporal_modality = gr.Dropdown(
            ["optical", "multispectral", "sar", "unknown"], value="optical", label="Modality"
        )
        temporal_cache = gr.Checkbox(value=True, label="Use cache")
        temporal_rerun = gr.Checkbox(value=False, label="Force rerun")
        temporal_submit = gr.Button("Run bi-temporal workflow", variant="primary")
        temporal_submit.click(
            bitemporal_api,
            [
                earlier_file,
                later_file,
                temporal_query,
                earlier_date,
                later_date,
                temporal_modality,
                temporal_cache,
                temporal_rerun,
            ],
            _outputs(),
            api_name="bitemporal",
        )

    with gr.Tab("Readiness"):
        readiness_output = gr.JSON(label="Health")
        contract_output = gr.JSON(label="Contract")
        gr.Button("Health (CPU only)").click(space_health, outputs=readiness_output, api_name="health")
        gr.Button("Contract (CPU only)").click(api_contract, outputs=contract_output, api_name="contract")


demo.queue(default_concurrency_limit=1, max_size=8)


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