#!/usr/bin/env python3
"""Gradio UI for query entity export on Hugging Face Spaces."""

from __future__ import annotations

import gradio as gr

from space_service import format_entry_json, get_service


def run_single_query(query: str) -> tuple[str, str]:
    query = (query or "").strip()
    if not query:
        return "", "Enter a clinical evidence query."

    try:
        entry = get_service().process_query(query)
    except FileNotFoundError as exc:
        return "", f"Service not ready: {exc}"
    except Exception as exc:  # noqa: BLE001
        return "", f"Error: {exc}"

    entities = entry.get("entities", [])
    summary = (
        f"Detected {len(entities)} entity/entities for query "
        f"({len(query)} characters)."
    )
    return format_entry_json(entry), summary


def run_split_export(
    split: str,
    limit: int,
    pretty: bool,
    progress=gr.Progress(track_tqdm=False),
) -> tuple[str, str | None]:
    if limit <= 0:
        return "Limit must be at least 1.", None

    try:
        svc = get_service()
    except FileNotFoundError as exc:
        return f"Service not ready: {exc}", None

    def _progress(done: int, total: int) -> None:
        progress(done / total, desc=f"Processing {done}/{total}")

    try:
        output_path, summary = svc.process_split(
            split,
            limit=limit,
            pretty=pretty,
            progress_callback=_progress,
        )
    except Exception as exc:  # noqa: BLE001
        return f"Error: {exc}", None

    return summary, str(output_path)


DESCRIPTION = """
Detect UMLS concepts in clinical query text with **QuickUMLS**, then map each CUI
to the best graph **AUI** using the bundled parquet cache.

**Modes**
- **Single query** — paste one evidence string and get entities immediately.
- **Dataset split** — run `train`, `val`, or `test` JSONL (bundled). Use a limit for
  quick checks; full test (~17k rows) can take 10–15 minutes after QuickUMLS loads.

First cold start downloads ~5 GB QuickUMLS data from S3 (configured via Space secrets).
"""

with gr.Blocks(title="Query Entity Export") as demo:
    gr.Markdown("# Query Entity Export")
    gr.Markdown(DESCRIPTION)

    with gr.Tab("Single query"):
        query_input = gr.Textbox(
            label="Clinical evidence / query text",
            placeholder="Type II Diabetes Mellitus Uncontrolled",
            lines=3,
        )
        single_btn = gr.Button("Detect entities", variant="primary")
        single_summary = gr.Textbox(label="Status", interactive=False)
        single_output = gr.Code(label="Result JSON", language="json")

        single_btn.click(
            fn=run_single_query,
            inputs=query_input,
            outputs=[single_output, single_summary],
        )
        gr.Examples(
            examples=[
                ["Type II Diabetes Mellitus Uncontrolled"],
                ["Mild intermittent asthma without complication"],
                ["Patient with acute kidney injury stage 3"],
            ],
            inputs=query_input,
        )

    with gr.Tab("Dataset split"):
        split_input = gr.Dropdown(
            choices=["test", "val", "train"],
            value="test",
            label="Split",
        )
        limit_input = gr.Number(
            value=10,
            precision=0,
            minimum=1,
            label="Row limit",
            info="Full test split has 17,243 rows (~10–15 min). Start small.",
        )
        pretty_input = gr.Checkbox(label="Pretty-print JSON", value=False)
        split_btn = gr.Button("Run export", variant="primary")
        split_status = gr.Textbox(label="Status", interactive=False)
        split_file = gr.File(label="Download results JSON")

        split_btn.click(
            fn=run_split_export,
            inputs=[split_input, limit_input, pretty_input],
            outputs=[split_status, split_file],
        )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)