"""
Gradio chat UI for the Hashie MedGemma deployment on Modal.

Connects to the deployed Modal streaming endpoint — no local GPU required.
Supports English and Swahili prompts.

Usage (local):
    uv run python modal-deploy/gradio_app.py
    uv run python modal-deploy/gradio_app.py --share

The same file is uploaded as app.py to Hugging Face Spaces by deploy_hf_space.py.
On HF Spaces, endpoint URLs are read from environment variables.
"""

import json
import os

import gradio as gr
import httpx

# ── Endpoint URLs (env vars for HF Spaces, defaults for local) ───────────────

STREAM_URL = os.environ.get(
    "STREAM_URL",
    "https://sb-modal-ws--hashie-medgemma-service-medgemmaservice-cha-894014.modal.run",
)
HEALTH_URL = os.environ.get(
    "HEALTH_URL",
    "https://sb-modal-ws--hashie-medgemma-service-medgemmaservice-health.modal.run",
)

TIMEOUT = httpx.Timeout(connect=300.0, read=300.0, write=30.0, pool=30.0)

# ── Constants ─────────────────────────────────────────────────────────────────

COUNTRIES = ["Ghana", "Uganda", "Kenya", "Ethiopia"]

EXAMPLES = [
    ["How would you describe sexually transmitted infections?"],
    ["Is HIV treatable?"],
    ["Can I access telehealth services for ongoing healthcare needs?"],
    ["What are the signs of trichomoniasis?"],
    ["Can oral sex with an HIV-positive person lead to infection?"],
    ["How do I calculate my menstrual cycle length when my period is irregular?"],
    ["Unawezaje kuelezea maambukizi ya zinaa?"],
    ["Je, VVU inaweza kutibika?"],
    ["Je, ninaweza kupata huduma za afya kwa njia ya simu kwa miadi ya ufuatiliaji?"],
    ["Dalili za ugonjwa wa trichomoniasis ni zipi?"],
    ["Je, ngono ya mdomo na mtu mwenye VVU inaweza kusababisha maambukizi?"],
    ["Ninawezaje kuhesabu urefu wa mzunguko wangu wa hedhi ikiwa hedhi yangu si ya kawaida?"],
]


# ── Streaming chat handler ────────────────────────────────────────────────────


def _extract_text(content) -> str:
    """Extract plain text from a Gradio message content field.

    Handles both plain strings (Gradio 5.x) and typed block lists
    like [{"type": "text", "text": "..."}] (Gradio 6.x).
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(
            block.get("text", "") for block in content if block.get("type") == "text"
        )
    return str(content)


def chat_stream(
    message,
    history: list,
    country: str,
    max_new_tokens: int,
    do_sample: bool,
    temperature: float,
):
    """Consume the deployed SSE endpoint and yield partial responses to Gradio."""

    user_text = (
        _extract_text(message)
        if isinstance(message, str)
        else _extract_text(message.get("text", message.get("content", "")))
    )

    payload = {
        "message": user_text,
        "country": country,
        "max_new_tokens": max_new_tokens,
        "do_sample": do_sample,
        "temperature": temperature,
    }

    partial = ""

    try:
        with httpx.stream(
            "POST", STREAM_URL, json=payload, timeout=TIMEOUT
        ) as resp:
            if resp.status_code != 200:
                resp.read()
                yield f"**Error {resp.status_code}:** {resp.text}"
                return

            for line in resp.iter_lines():
                if not line.startswith("data: "):
                    continue

                data_str = line[6:]
                if data_str == "[DONE]":
                    break

                try:
                    token = json.loads(data_str)["token"]
                    partial += token
                    yield partial
                except (json.JSONDecodeError, KeyError):
                    pass

    except httpx.ConnectError:
        yield (
            "**Connection failed.** The service may be scaling up from zero. "
            "Please wait a minute and try again."
        )
    except httpx.ReadTimeout:
        yield partial + "\n\n**[Response timed out]**"


# ── Gradio UI (built at module level so HF Spaces auto-discovers it) ─────────

with gr.Blocks(title="Hashie — SRH Medical Assistant") as demo:
    with gr.Row():
        # ── Left sidebar: settings ────────────────────────────────────
        with gr.Column(scale=1, min_width=260):
            gr.Markdown("## Settings")

            country_dropdown = gr.Dropdown(
                choices=COUNTRIES,
                value="Ghana",
                label="Country",
                info="Tailors responses to the selected country.",
            )

            with gr.Accordion("Generation", open=False):
                max_tokens_slider = gr.Slider(
                    minimum=64,
                    maximum=1024,
                    value=512,
                    step=64,
                    label="Max new tokens",
                )
                do_sample_checkbox = gr.Checkbox(
                    value=False,
                    label="Sampling (off = greedy)",
                    info="Enable for more varied responses.",
                )
                temperature_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.3,
                    step=0.05,
                    label="Temperature",
                    info="Only active when Sampling is on.",
                )

        # ── Main area: chat ───────────────────────────────────────────
        with gr.Column(scale=4):
            gr.ChatInterface(
                fn=chat_stream,
                additional_inputs=[
                    country_dropdown,
                    max_tokens_slider,
                    do_sample_checkbox,
                    temperature_slider,
                ],
                title="Hashie — SRH Medical Assistant (MedGemma 4B)",
                description=(
                    "Fine-tuned MedGemma 4B for sexual and reproductive health Q&A. "
                    "Supports **English** and **Swahili** — just type in either language. "
                    "Select a country in the left panel to tailor responses, "
                    "then ask any SRH question.\n\n"
                    "*Powered by a serverless GPU backend on Modal — "
                    "the first message after idle may take a minute while the model loads.*"
                ),
                examples=EXAMPLES,
                cache_examples=False,
                autofocus=True,
            )

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Hashie Gradio Chat UI")
    parser.add_argument(
        "--share", action="store_true", help="Create a public Gradio share link"
    )
    args = parser.parse_args()

    demo.launch(share=args.share)