"""ZeroGPU Gradio Space — crisis triage student model demo, with voice in/out
and a model switcher between trained adapters.

ZeroGPU only attaches a real GPU for the duration of a function decorated
with @spaces.GPU -- at module/import time there is no GPU visible, so every
model must load on CPU first and move to "cuda" only inside that function.

Models are LAZY-LOADED: nothing but the dropdown's default entry loads at
startup. Switching the dropdown loads that model on CPU (cached after first
use) and frees the previous one, so you never hold a 1.5B and a 7B model in
CPU RAM at the same time. STT/Whisper and TTS/VITS are always-on and shared
across model choices.

Pipeline per click (all inside ONE @spaces.GPU call, to use a single GPU
borrow instead of three):
    voice in  -> Whisper STT  -> transcript
    transcript -> triage adapter (selected model) -> structured JSON
    JSON["plain_language_summary"] -> TTS -> spoken response

EDIT this dict to add/remove selectable models:
    MODEL_REGISTRY
"""
import spaces  # must be imported before torch/transformers/peft on ZeroGPU

import gc
import html
import json
import re

import gradio as gr
import numpy as np
import torch
from peft import PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    VitsModel,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

MODEL_REGISTRY = {
    "Qwen2.5-1.5B (fast)": {
        "base": "Qwen/Qwen2.5-1.5B-Instruct",
        "adapter": "NathanPereira/crisis-triage-adapter",
    },
    "Qwen3-4B (balanced)": {
        "base": "Qwen/Qwen3-4B",
        "adapter": "NathanPereira/crisis-triage-adapter-qwen3-4b",
    },
    "DeepSeek-R1-Distill-Qwen-7B (stronger)": {
        "base": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
        "adapter": "NathanPereira/crisis-triage-adapter-r1-distill-7b",
    },
}
DEFAULT_MODEL_NAME = "Qwen2.5-1.5B (fast)"

ASR_MODEL_ID = "openai/whisper-small"
TTS_MODEL_ID = "facebook/mms-tts-eng"

SYSTEM_PROMPT = (
    "You are a research-only crisis triage assistant. "
    "Return only valid JSON matching the triage schema. "
    "Use exact evidence spans copied from the input text. "
    "Do not diagnose, prescribe medication, or provide therapy instructions."
)

# name -> (tokenizer, model) on CPU; populated lazily, at most one entry at a time
_loaded = {}


def _get_triage_model(name: str):
    if name in _loaded:
        return _loaded[name]

    # evict whatever was loaded before so we never hold two student models in RAM
    if _loaded:
        old_name = next(iter(_loaded))
        del _loaded[old_name]
        gc.collect()

    cfg = MODEL_REGISTRY[name]
    print(f"loading triage model '{name}' (CPU) ...")
    tok = AutoTokenizer.from_pretrained(cfg["base"])
    base = AutoModelForCausalLM.from_pretrained(cfg["base"], torch_dtype=torch.float16)
    model = PeftModel.from_pretrained(base, cfg["adapter"], torch_device="cpu").eval()
    _loaded[name] = (tok, model)
    return tok, model


print("loading Whisper STT (CPU) ...")
asr_processor = WhisperProcessor.from_pretrained(ASR_MODEL_ID)
asr_model = WhisperForConditionalGeneration.from_pretrained(ASR_MODEL_ID, torch_dtype=torch.float16).eval()

print("loading TTS (CPU) ...")
tts_tok = AutoTokenizer.from_pretrained(TTS_MODEL_ID)
tts_model = VitsModel.from_pretrained(TTS_MODEL_ID).eval()

print(f"warming default triage model: {DEFAULT_MODEL_NAME} ...")
_get_triage_model(DEFAULT_MODEL_NAME)
print("all models ready")


def _salvage_tier(raw: str) -> int | None:
    m = re.search(r'"risk_tier"\s*:\s*([0-3])', raw)
    return int(m.group(1)) if m else None


def _extract_json(raw: str) -> dict:
    if "</think>" in raw:
        raw = raw.split("</think>", 1)[1]
    s, e = raw.find("{"), raw.rfind("}")
    if s < 0 or e <= s:
        raise ValueError("no JSON object")
    return json.loads(raw[s : e + 1])


def _transcribe(audio) -> str:
    # gradio Audio(type="numpy") gives (sample_rate, np.ndarray); no file I/O, no librosa.
    import torchaudio

    sr, data = audio
    wav = torch.as_tensor(data).float()
    if wav.ndim == 2:                      # stereo -> mono
        wav = wav.mean(dim=1)
    if data.dtype.kind in "iu":            # int PCM -> [-1, 1]
        wav = wav / float(np.iinfo(data.dtype).max)
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    inputs = asr_processor(wav.numpy(), sampling_rate=16000, return_tensors="pt")
    input_features = inputs.input_features.to("cuda", dtype=torch.float16)
    with torch.no_grad():
        pred_ids = asr_model.generate(input_features, language="en", task="transcribe")
    return asr_processor.batch_decode(pred_ids, skip_special_tokens=True)[0].strip()


def _normalize_tier(tier):
    """Collapse the many risk_tier shapes the model emits into an int 0..3.

    Handles ints, floats, numeric strings ("2"), and free-text labels
    ("High Risk", "imminent", "low") so downstream phrasing is consistent.
    Returns None if nothing recognizable.
    """
    if isinstance(tier, bool) or tier is None:
        return None
    if isinstance(tier, (int, float)):
        t = int(tier)
        return t if 0 <= t <= 3 else None
    if isinstance(tier, str):
        s = tier.strip().lower()
        m = re.search(r"[0-3]", s)
        if m:
            return int(m.group())
        if any(w in s for w in ("imminent", "critical", "emergency", "severe", "tier 3")):
            return 3
        if "high" in s:
            return 2
        if any(w in s for w in ("moderate", "medium", "elevated")):
            return 1
        if any(w in s for w in ("low", "minimal", "none", "no risk")):
            return 0
    return None


def _first_text(value) -> str:
    """Pull a usable sentence out of a str / list / dict field."""
    if isinstance(value, str):
        return value.strip()
    if isinstance(value, (list, tuple)):
        for v in value:
            t = _first_text(v)
            if t:
                return t
    if isinstance(value, dict):
        for v in value.values():
            t = _first_text(v)
            if t:
                return t
    return ""


# tier-aware openers, written to be spoken aloud like a person talking, not a report.
# each has a connector for folding a next-step in naturally instead of a labeled clause.
_TIER_VOICE = {
    0: dict(
        opener="Thanks for sharing that with me. From what you've described, things "
               "don't sound like they're at a crisis point right now.",
        connector="That said, it can still help to",
    ),
    1: dict(
        opener="I can hear that you're going through a tough time. What you're "
               "describing sounds like real distress, and it's worth taking seriously.",
        connector="One thing that might help is to",
    ),
    2: dict(
        opener="I'm genuinely concerned about what you've shared. This sounds serious, "
               "and I want you to know your safety matters right now.",
        connector="Please",
    ),
    3: dict(
        opener="What you've told me sounds like an emergency, and I'm worried about "
               "your safety right now.",
        connector="Right now, please",
    ),
}
_DEFAULT_VOICE = dict(
    opener="Thank you for trusting me with this. Whatever you're going through, "
           "you don't have to face it alone.",
    connector="It could help to",
)
# always-on safety line for higher tiers, regardless of what JSON fields exist
_SAFETY_LINE = ("And if you're ever in immediate danger, please call your local "
                "emergency number, or reach out to a crisis line right away.")


def _lowercase_lead(s: str) -> str:
    """Lowercase a sentence's first letter so it reads naturally mid-sentence."""
    return s[:1].lower() + s[1:] if s else s


def _strip_trailing_period(s: str) -> str:
    return s[:-1] if s.endswith(".") else s


def _build_summary(parsed: dict) -> str:
    """Compose a warm, conversational, tier-aware spoken triage summary.

    Written for a synthetic chain-of-thought triage adapter: it talks to the person
    the way a caring listener would, folds in one concrete next step as part of the
    same flowing sentence rather than a labeled clause, and always closes high-risk
    tiers with a safety line -- even when the model drifts off-schema.
    Prefers the trained `plain_language_summary` verbatim when present.
    """
    if not isinstance(parsed, dict):
        return ""

    direct = parsed.get("plain_language_summary")
    if isinstance(direct, str) and direct.strip():
        return direct.strip()

    tier = _normalize_tier(parsed.get("risk_tier"))
    voice = _TIER_VOICE.get(tier, _DEFAULT_VOICE)
    sentence = voice["opener"]

    # fold one concrete next step into the same breath as the opener, instead of
    # a separate labeled sentence
    for key in ("recommended_next_step", "action_steps", "resources",
                "clinical_rationale", "risk_factors"):
        action = _first_text(parsed.get(key))
        if action:
            action = _strip_trailing_period(_lowercase_lead(action.strip()))
            sentence = f"{sentence} {voice['connector']} {action}."
            break

    escalate = (parsed.get("escalation_required") in (True, "true", "True", "yes", 1)
                or (tier is not None and tier >= 2))
    if escalate:
        sentence = f"{sentence} {_SAFETY_LINE}"

    return sentence.strip()


# badge color/label per normalized tier, used to render the report card
_TIER_META = {
    0: dict(color="#2e7d32", bg="#e8f5e9", label="Tier 0 · Minimal"),
    1: dict(color="#a16207", bg="#fef9c3", label="Tier 1 · Moderate"),
    2: dict(color="#c2410c", bg="#ffedd5", label="Tier 2 · High"),
    3: dict(color="#b91c1c", bg="#fee2e2", label="Tier 3 · Imminent"),
}
_TIER_META_UNKNOWN = dict(color="#374151", bg="#f3f4f6", label="Tier unclear")

# keys rendered up front, in this order, with friendly labels + how to render them
_REPORT_FIELDS = [
    ("evidence_spans", "Evidence from the text", "quotes"),
    ("risk_factors", "Risk factors", "list"),
    ("protective_factors", "Protective factors", "list"),
    ("cssrs_axes", "C-SSRS axes", "table"),
    ("clinical_rationale", "Clinical rationale", "text"),
    ("recommended_next_step", "Recommended next step", "highlight"),
    ("uncertainty_flags", "Uncertainty flags", "list"),
]
# fields shown elsewhere (tier badge, escalation banner) or internal -- skip in the body
_REPORT_SKIP_KEYS = {
    "risk_tier", "plain_language_summary", "escalation_required",
    "_parse_error", "_raw", "confidence",
}


def _esc(value) -> str:
    return html.escape(str(value), quote=True)


def _render_inline(value) -> str:
    """Render a str/list/dict value as a flowing HTML fragment (no own heading)."""
    if isinstance(value, str):
        return _esc(value)
    if isinstance(value, (list, tuple)):
        items = [_render_inline(v) for v in value if _render_inline(v)]
        if not items:
            return ""
        return "<ul class='ct-list'>" + "".join(f"<li>{i}</li>" for i in items) + "</ul>"
    if isinstance(value, dict):
        rows = "".join(
            f"<tr><td class='ct-k'>{_esc(k)}</td><td>{_render_inline(v)}</td></tr>"
            for k, v in value.items() if _render_inline(v)
        )
        return f"<table class='ct-table'>{rows}</table>" if rows else ""
    if value in (None, "", [], {}):
        return ""
    return _esc(value)


def _render_section(label: str, value, kind: str) -> str:
    if value in (None, "", [], {}):
        return ""
    if kind == "quotes" and isinstance(value, (list, tuple)):
        body = "".join(f"<blockquote class='ct-quote'>“{_esc(v)}”</blockquote>" for v in value if str(v).strip())
    elif kind == "highlight":
        body = f"<div class='ct-highlight'>{_render_inline(value)}</div>"
    else:
        body = _render_inline(value)
    if not body:
        return ""
    return f"<div class='ct-section'><div class='ct-section-title'>{_esc(label)}</div>{body}</div>"


def render_report(parsed: dict) -> str:
    """Render the triage JSON as a readable HTML report card instead of raw JSON.

    Walks the known schema fields first (evidence, rationale, next step, ...) in a
    fixed friendly order, then falls back to rendering any extra/unexpected keys the
    model produced (schema drift), so nothing the model says is silently dropped.
    """
    if not isinstance(parsed, dict):
        return "<div class='ct-card'><i>No structured output to show.</i></div>"

    if parsed.get("_parse_error"):
        return (
            "<div class='ct-card'><div class='ct-section'>"
            "<div class='ct-section-title'>Couldn't parse a structured response</div>"
            "<p>The model's raw output didn't contain valid JSON -- see the raw output box below.</p>"
            "</div></div>"
        )

    tier = _normalize_tier(parsed.get("risk_tier"))
    meta = _TIER_META.get(tier, _TIER_META_UNKNOWN)
    confidence = parsed.get("confidence")
    conf_html = ""
    if isinstance(confidence, (int, float)):
        conf_html = f"<span class='ct-confidence'>confidence {confidence:.0%}" \
                     f"</span>" if confidence <= 1 else f"<span class='ct-confidence'>confidence {confidence}</span>"

    escalate = (parsed.get("escalation_required") in (True, "true", "True", "yes", 1)
                or (tier is not None and tier >= 2))

    sections = "".join(
        _render_section(label, parsed.get(key), kind) for key, label, kind in _REPORT_FIELDS
    )

    extra_keys = [k for k in parsed if k not in _REPORT_SKIP_KEYS
                  and k not in {f[0] for f in _REPORT_FIELDS}]
    extras = "".join(
        _render_section(k.replace("_", " ").title(), parsed.get(k), "text") for k in extra_keys
    )

    banner = (
        "<div class='ct-escalate'>⚠ This case may need urgent attention or escalation.</div>"
        if escalate else ""
    )

    return f"""
    <div class="ct-card">
      <style>
        .ct-card {{ font-family: inherit; }}
        .ct-badge {{ display:inline-block; padding:4px 12px; border-radius:999px;
                     font-weight:600; font-size:0.85em; color:{meta['color']};
                     background:{meta['bg']}; margin-right:8px; }}
        .ct-confidence {{ font-size:0.85em; opacity:0.75; }}
        .ct-escalate {{ margin-top:10px; padding:8px 12px; border-radius:8px;
                        background:#fee2e2; color:#b91c1c; font-weight:600; font-size:0.9em; }}
        .ct-section {{ margin-top:16px; }}
        .ct-section-title {{ font-weight:600; font-size:0.85em; text-transform:uppercase;
                              letter-spacing:0.03em; opacity:0.6; margin-bottom:6px; }}
        .ct-list {{ margin:0; padding-left:1.2em; }}
        .ct-quote {{ margin:4px 0; padding:6px 10px; border-left:3px solid #d1d5db;
                     font-style:italic; opacity:0.85; }}
        .ct-highlight {{ padding:10px 12px; border-radius:8px; background:#eff6ff;
                          border:1px solid #bfdbfe; }}
        .ct-table {{ width:100%; border-collapse:collapse; }}
        .ct-table td {{ padding:4px 8px; vertical-align:top; border-bottom:1px solid #e5e7eb; }}
        .ct-table td.ct-k {{ font-weight:600; white-space:nowrap; opacity:0.7; }}
      </style>
      <div><span class="ct-badge">{_esc(meta['label'])}</span>{conf_html}</div>
      {banner}
      {sections}
      {extras}
    </div>
    """


def _synthesize(text: str) -> tuple[int, np.ndarray]:
    inputs = tts_tok(text, return_tensors="pt").to("cuda")
    with torch.no_grad():
        waveform = tts_model(**inputs).waveform
    audio = waveform.squeeze().to(torch.float32).cpu().numpy()
    return tts_model.config.sampling_rate, audio


@spaces.GPU  # borrows a real GPU only for the duration of this call
def analyze(model_name: str, text: str, audio) -> tuple[str, str, str, str, str, tuple]:
    tok, model = _get_triage_model(model_name)

    asr_model.to("cuda")
    tts_model.to("cuda")
    model.to("cuda")

    transcript = ""
    if audio is not None:
        try:
            transcript = _transcribe(audio)
            text = transcript or text
        except Exception as exc:
            transcript = f"(speech-to-text failed: {exc})"

    if not text or not text.strip():
        placeholder = "<div class='ct-card'><i>Enter text or record audio first.</i></div>"
        return placeholder, "", "", transcript, "", None

    messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": text}]
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tok(prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        out = model.generate(
            **inputs, max_new_tokens=400, do_sample=False, pad_token_id=tok.eos_token_id
        )
    raw = tok.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True)

    try:
        parsed = _extract_json(raw)
    except Exception:
        parsed = {"risk_tier": _salvage_tier(raw), "_parse_error": True, "_raw": raw}

    summary = _build_summary(parsed) or (
        "I couldn't generate a clear summary -- please check the structured output "
        "or contact a crisis line directly."
    )
    try:
        speech_audio = _synthesize(summary)
    except Exception as exc:
        summary = f"{summary}\n\n(text-to-speech failed: {exc})"
        speech_audio = None

    return (
        render_report(parsed),
        json.dumps(parsed, indent=2, ensure_ascii=False),
        raw,
        transcript,
        summary,
        speech_audio,
    )


with gr.Blocks(title="Crisis Triage — Research Demo") as demo:
    gr.Markdown(
        "## Crisis Triage — Research Demo\n"
        "Running on **ZeroGPU**, with voice input (Whisper STT), voice output (TTS), "
        "and a switchable student model.\n\n"
        "⚠ Research prototype. Not a clinical tool. "
        "Outputs are not a substitute for professional judgment or emergency services. "
        "If you or someone else is in immediate danger, contact local emergency services "
        "or a crisis line directly."
    )
    model_choice = gr.Dropdown(
        label="Model",
        choices=list(MODEL_REGISTRY.keys()),
        value=DEFAULT_MODEL_NAME,
        info="Switching loads the new adapter on first use (a few seconds) and evicts the previous one.",
    )
    with gr.Row():
        inp = gr.Textbox(label="Text to assess", lines=5, placeholder="Type a message or transcript…")
        mic = gr.Audio(label="...or record/upload voice", sources=["microphone", "upload"], type="numpy")
    btn = gr.Button("Analyze", variant="primary")

    transcript_box = gr.Textbox(label="Transcribed speech (if audio was used)", lines=2, interactive=False)
    out_report = gr.HTML(label="Triage report")
    with gr.Accordion("Raw model output (JSON / debug)", open=False):
        out_json = gr.Code(label="Structured JSON output", language="json")
        out_raw = gr.Textbox(label="Raw model output", lines=4)
    out_summary = gr.Textbox(label="Plain-language summary (spoken aloud)", lines=2, interactive=False)
    out_audio = gr.Audio(label="Spoken response", type="numpy")

    btn.click(
        analyze,
        inputs=[model_choice, inp, mic],
        outputs=[out_report, out_json, out_raw, transcript_box, out_summary, out_audio],
    )

demo.launch()