"""Thyroid Initial Visit — Endocrine / Family Medicine outpatient.

Form-based Gradio app. Calls Claude once per generation, passing structured
form data (phenotype tags + ROS + PE + prior data) plus any lab/imaging
screenshots. The model returns a 4-section output:

  1) Red-flag alert (only if detected)
  2) History + ROS narrative
  3) Physical examination findings
  4) Initial workup orders (checklist + narrative)

Tab 2 is a pure-Python ACR TI-RADS calculator (no API call).
"""

from __future__ import annotations

import base64
import mimetypes
import os
import re
import tempfile
from datetime import datetime, timedelta, timezone
from pathlib import Path

TAIPEI_TZ = timezone(timedelta(hours=8))


def _patch_gradio_client_schema_bug() -> None:
    """gradio_client/utils.py crashes when a JSON schema value is `bool` instead
    of `dict` (newer pydantic emits `additionalProperties: False`). Patch so the
    schema walker treats booleans as `Any`. Must run before gradio is imported.
    """
    import gradio_client.utils as gcu

    orig_get_type = gcu.get_type
    orig_walk = gcu._json_schema_to_python_type

    def _get_type(schema):
        if not isinstance(schema, dict):
            return "Any"
        return orig_get_type(schema)

    def _walk(schema, defs):
        if isinstance(schema, bool):
            return "Any"
        return orig_walk(schema, defs)

    gcu.get_type = _get_type
    gcu._json_schema_to_python_type = _walk


_patch_gradio_client_schema_bug()

import anthropic
import gradio as gr

import gdrive

# ──────────────────────────────────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────────────────────────────────

MODEL = os.environ.get("THYROID_MODEL", "claude-sonnet-4-6")
PASSWORD = os.environ.get("THYROID_PASSWORD", "").strip()

SYSTEM_PROMPT = """You are a clinical scribe assisting an endocrinology / family-medicine outpatient physician at a Taiwanese hospital. You produce English thyroid initial-visit notes in a fixed 4-section structure for outpatient chart documentation.

ABSOLUTE RULES — these override any instinct to be helpful or thorough:

1. **No inference of clinical facts.** Use ONLY what the physician documented in the form or what is directly visible in uploaded screenshots. Never invent symptoms, history, lab values, drug doses, or diagnoses. Never extrapolate "likely" findings.
2. **Missing data** → write `Unknown/Not documented`. Never guess.
3. **Blank fields or fields marked `nil`, `(nil)`, `none`, `n/a`, `na`, `skip`, or obvious empty placeholders** → omit that section/sentence entirely. Do not write `Unknown/Not documented` for these — just skip.
4. **Concise but complete.** Include only what the physician explicitly entered. Do not list every pertinent negative.
5. **Phenotype-driven order selection.** The physician checks one or more phenotype tags (hyperthyroid / hypothyroid / goiter-nodule / abnormal-lab). Workup orders in Section 4 are selected per the rules below, NOT a blanket panel.
6. **Lab screenshots:** extract only values you can clearly read. If illegible, OMIT — do not guess. State units as printed.
7. **Section 1 (Red flags) is conditional.** Only output Section 1 when one of the explicit red-flag patterns below is present in the documented data. Otherwise OMIT Section 1 entirely.
8. **You are a scribe + suggestion engine, NOT a diagnostician.** Workup suggestions are explicitly framed as "suggested orders for physician review." Section 4D differential is one short paragraph of pattern recognition, not a final diagnosis.

OUTPUT FORMAT — produce EXACTLY this structure, no preamble, no closing remarks. Section headers verbatim. If a section is OMITTED per rules, skip its heading entirely (do not leave an empty heading).

```
🚨 RED FLAGS — urgent attention required
- {flag name}: {one-sentence reason from documented data} → {recommended action}
(only if at least one red-flag pattern detected)

1) History and Review of Systems

Phenotype tags (per physician): {comma-separated tags physician checked, or "not specified — inferred from notes" if none}

This is a {age}-year-old {man/woman} who presents to the {clinic type} for {chief concern}, with symptom onset {onset / time course}.

According to the patient, {HPI narrative organized by checked phenotype(s): hyperthyroid symptoms, hypothyroid symptoms, neck/compressive symptoms, eye symptoms, pain/inflammatory symptoms, abnormal-lab discovery — only include sub-narratives for documented categories}. Pertinent negatives include {only those explicitly entered}.

Past medical history is significant for {documented PMH}. {Sentence on prior thyroid disease/surgery/RAI if entered.} {Family history sentence if entered.}

Medication and exposure history: {documented meds + exposures}. Note specifically: {flag any of: levothyroxine, antithyroid drugs, amiodarone, lithium, immune-checkpoint inhibitor, TKI, biotin, iodine contrast/kelp, head-neck radiation, prior RAI — IF entered}.

{For women only, if documented: pregnancy / postpartum / breastfeeding / menstrual / fertility status.}

2) Physical examination

Vital signs: {documented vitals}.
General: {general appearance}.
{Eye examination: only if eye findings documented.}
{Neck inspection / palpation / auscultation: thyroid size, consistency, nodule (location, size, mobility, tenderness), tracheal deviation, bruit — only documented findings.}
{Cervical lymph nodes: documented findings.}
{Hand and neuromuscular: tremor, palm warmth/moisture, deep tendon reflex, proximal strength.}
{Cardiovascular: rhythm, rate, murmur.}
{Skin and hair: warm/moist vs dry/coarse, pretibial myxedema, hair pattern.}
{Other documented findings.}

3) Prior data on file (only if documented)

{One short paragraph summarizing physician's free-text prior labs / prior imaging / prior FNA Bethesda category, IF entered. If all three are blank, OMIT Section 3 entirely.}

4) Initial workup plan

A) Suggested orders (checklist for physician review)

Select orders based on checked phenotype(s). For each suggested order, append a 3-8 word reason. Use ☐ checkbox format.

Phenotype-driven core panel:
- Always: ☐ TSH ☐ Free T4
- Hyperthyroid suspected: ☐ Total T3 (or Free T3) ☐ TRAb (or TSI) ☐ CBC with differential — baseline before ATD ☐ AST/ALT, total bilirubin — baseline before ATD ☐ ECG — palpitation/AF check
- Hypothyroid suspected: ☐ Anti-TPO Ab ☐ ± Anti-Tg Ab ☐ Lipid profile ☐ CBC ☐ Na, Creatinine
- Goiter/nodule: ☐ Thyroid ultrasound ☐ Cervical lymph node ultrasound ☐ Thyroid scan / RAIU — only if TSH suppressed and nodule present
- Abnormal-lab only: ☐ Repeat TSH + Free T4 in 4-6 weeks ☐ ± Anti-TPO if TSH ↑ ☐ ± TRAb if TSH ↓

Conditional additions:
- Pregnant or planning pregnancy → ☐ Pregnancy test ☐ Adjust TSH target per ATA pregnancy guideline
- Female of reproductive age + likely ATD plan → ☐ Pregnancy test
- Subacute thyroiditis suspected (neck pain + transient hyper) → ☐ ESR ☐ CRP
- Periodic paralysis suspected (young male, weakness, hyper) → ☐ K, Mg, Phosphate, CK
- Severe hypothyroid + fatigue + anemia → ☐ B12, Ferritin
- Long-standing thyroid mass / firm fixed mass / family h/o MTC → ☐ Calcitonin (NOT routine for nodule screening)

B) Order narrative (single-line, for direct paste into HIS Plan section)

Plan: {comma-separated list of the checked orders above, in the order listed}. {If thyroid US planned: "Thyroid ultrasound with cervical lymph node survey."} {If RAIU planned: "Radioactive iodine uptake / thyroid scan when TSH confirmed suppressed."} Follow-up in {2-4 weeks for hyper / 4-6 weeks for hypo / 6-8 weeks for asymptomatic abnormal-lab / 2-4 weeks for nodule after US} pending results.

C) Patient instructions (1-2 short sentences)
{e.g. "Fast not required for TSH/FT4. If on biotin supplement, stop ≥48h before draw." OR "Avoid iodine contrast and kelp supplements before RAIU scan." Choose only relevant ones.}

D) Differential and assessment thinking (brief, 2-4 sentences)
{Synthesize: which etiologies to keep on the differential given the phenotype and documented findings. Examples: "Pattern is consistent with overt hyperthyroidism; differential includes Graves disease (most likely given diffuse goiter + eye signs if present), toxic multinodular goiter, toxic adenoma, and thyroiditis. TRAb and Doppler ultrasound will help distinguish." Keep neutral if data sparse.}
```

RED-FLAG DETECTION RULES — output Section 1 ONLY when one or more triggered:

A. **Thyroid storm risk**: documented hyperthyroid context (clinical hyper symptoms OR known Graves OR prior hyper labs) PLUS any TWO of: fever (BT > 38.5°C), HR > 130, atrial fibrillation, altered mental status, vomiting/diarrhea, jaundice.
→ Action: "Immediate ED transfer for thyroid storm workup (Burch-Wartofsky score)."

B. **Myxedema coma risk**: documented severe hypothyroid context PLUS any TWO of: altered mental status, hypothermia (BT < 35°C), bradycardia (HR < 50), hypotension (SBP < 90), severe hyponatremia, hypoventilation.
→ Action: "Immediate ED transfer for myxedema coma workup; ICU-level care anticipated."

C. **Airway compromise from thyroid mass**: any one of: stridor, severe dyspnea, orthopnea attributed to neck mass, rapid enlargement of mass, positive Pemberton sign.
→ Action: "Urgent ENT consult and chest CT for tracheal compression / retrosternal extension."

D. **High-suspicion thyroid malignancy on exam**: fixed hard thyroid mass + any one of (hoarseness, fixed cervical lymph node, h/o childhood head-neck radiation, family h/o MTC/MEN2).
→ Action: "Expedited thyroid ultrasound within 1-2 weeks; consider direct laryngoscopy if hoarseness present."

E. **Graves orbitopathy sight-threatening features**: documented Graves context PLUS any one of (visual acuity decline, color vision change, severe eye pain, corneal exposure / inability to close lid).
→ Action: "Urgent ophthalmology referral for dysthyroid optic neuropathy / corneal protection."

If MULTIPLE red flags detected, list each as a separate bullet under the Section 1 header.

OUTPUT LANGUAGE: English only. Do not include Chinese in the output. Do not add disclaimers, AI notes, or "based on the information provided" preambles.
"""


# ──────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────


def _is_nil(s: str | None) -> bool:
    if s is None:
        return True
    t = s.strip().lower()
    return t in {"", "nil", "(nil)", "none", "n/a", "na", "skip", "-", "--"}


def _line(label: str, value: str | None) -> str:
    if _is_nil(value):
        return f"- {label}: (not documented / skipped)"
    return f"- {label}: {value.strip()}"


def _section(title: str, lines: list[str], skip: bool = False) -> str:
    if skip:
        return f"## {title}\n(physician marked this section as nil — OMIT from output)\n"
    return f"## {title}\n" + "\n".join(lines) + "\n"


def _image_block(filepath: str) -> dict:
    mime, _ = mimetypes.guess_type(filepath)
    if not mime or not mime.startswith("image/"):
        mime = "image/png"
    data = base64.standard_b64encode(Path(filepath).read_bytes()).decode("utf-8")
    return {
        "type": "image",
        "source": {"type": "base64", "media_type": mime, "data": data},
    }


# ──────────────────────────────────────────────────────────────────────────
# Form → structured prompt
# ──────────────────────────────────────────────────────────────────────────


def _phenotype_summary(form: dict) -> str:
    tags = []
    if form.get("pt_hyper"):
        tags.append("hyperthyroid")
    if form.get("pt_hypo"):
        tags.append("hypothyroid")
    if form.get("pt_nodule"):
        tags.append("goiter / nodule")
    if form.get("pt_abnormal_lab"):
        tags.append("abnormal-lab discovery")
    return ", ".join(tags) if tags else "not specified — physician left all four boxes unchecked"


def build_user_message(form: dict, screenshots: list[str]) -> list[dict]:
    parts: list[str] = ["# Physician-entered thyroid initial-visit data\n"]

    parts.append(_section("Visit context", [
        _line("Age", form["age"]),
        _line("Sex", form["sex"]),
        _line("Clinic type", form["clinic_type"]),
        _line("Visit date/time", form["visit_datetime"]),
        _line("Chief concern", form["chief_concern"]),
        _line("Symptom onset / time course", form["onset"]),
    ]))

    parts.append(_section("Phenotype tags (checked by physician)", [
        f"- Tags: {_phenotype_summary(form)}",
        "- Use these to drive Section 4A order selection. If none checked, infer best-fit phenotype from documented HPI/PE and STATE that inference at top of Section 1.",
    ]))

    parts.append(_section("Review of systems — hyperthyroid axis", [
        _line("Hyper symptoms (palpitation/tremor/heat intolerance/sweating/weight loss/diarrhea/insomnia/anxiety/proximal weakness/menstrual)", form["ros_hyper"]),
    ]))

    parts.append(_section("Review of systems — hypothyroid axis", [
        _line("Hypo symptoms (fatigue/cold intolerance/weight gain/constipation/dry skin/hair loss/edema/hoarseness/cognitive slowing/menorrhagia)", form["ros_hypo"]),
    ]))

    parts.append(_section("Review of systems — neck / compressive", [
        _line("Neck symptoms (mass, pain, dysphagia, dyspnea, orthopnea, hoarseness, rapid enlargement, choking, Pemberton)", form["ros_neck"]),
    ]))

    parts.append(_section("Review of systems — eye (Graves)", [
        _line("Eye symptoms (proptosis, dry eye, diplopia, eye pain, visual decline, photophobia, lid lag)", form["ros_eye"]),
    ]))

    parts.append(_section("Review of systems — pain / inflammatory", [
        _line("Pain pattern (thyroid tender, neck pain radiating to jaw/ear, recent URI, fever, postpartum status, transient hyper→hypo course)", form["ros_pain"]),
    ]))

    parts.append(_section("Past medical history (thyroid + related)", [
        _line("Prior thyroid disease (Graves/Hashimoto/nodule/cancer/thyroiditis/surgery/RAI)", form["pmh_thyroid"]),
        _line("Other autoimmune disease (T1DM/vitiligo/celiac/RA/SLE/Addison/pernicious anemia)", form["pmh_autoimmune"]),
        _line("Cardiovascular (AF/HF/CAD/HTN/osteoporosis)", form["pmh_cv"]),
        _line("Psychiatric / other relevant (anxiety/depression/liver/renal/leukopenia/recent malignancy)", form["pmh_other"]),
    ]))

    parts.append(_section("Family history", [
        _line("Thyroid disease in family (Graves/Hashimoto/nodule/thyroid cancer/MEN2/MTC)", form["fh_thyroid"]),
        _line("Other autoimmune family history", form["fh_autoimmune"]),
    ]))

    parts.append(_section("Medications, supplements, exposures", [
        _line("Thyroid-specific drugs (levothyroxine dose/timing/co-meds, methimazole, PTU, T3 prep)", form["meds_thyroid"]),
        _line("Other thyroid-relevant drugs (amiodarone, lithium, IFN-α, immune checkpoint inhibitor, TKI, biotin, OCP/estrogen, anti-seizure)", form["meds_other_thyroid"]),
        _line("Iodine exposure (recent CT/cath contrast, kelp/seaweed/kelp supplement, betadine, Chinese-medicine herbs)", form["exposure_iodine"]),
        _line("Radiation exposure (childhood/young-adult head-neck RT)", form["exposure_radiation"]),
    ]))

    parts.append(_section("Female-specific (omit entirely if not applicable / male)", [
        _line("Pregnancy / planning / breastfeeding / postpartum (specify months postpartum)", form["female_pregnant"]),
        _line("Menstrual / fertility", form["female_menses"]),
    ]))

    parts.append(_section("Physical examination", [
        _line("Vital signs (BT °C / HR bpm / RR / BP mmHg / SpO2 / body weight kg)", form["pe_vitals"]),
        _line("General appearance", form["pe_general"]),
        _line("Eye exam (lid retraction, lid lag, proptosis, EOM, visual acuity, conjunctival injection, periorbital edema)", form["pe_eye"]),
        _line("Neck inspection (visible enlargement, asymmetry, surgical scar, venous engorgement, moves with swallow)", form["pe_neck_inspect"]),
        _line("Neck palpation (size, consistency soft/firm/hard, surface smooth/nodular, tenderness, mobility, nodule location/size/border, tracheal deviation, retrosternal extension)", form["pe_neck_palp"]),
        _line("Neck auscultation (thyroid bruit)", form["pe_neck_ausc"]),
        _line("Cervical lymph nodes (location, size, consistency, mobility — note any fixed/hard/cystic/supraclavicular)", form["pe_lymph"]),
        _line("Hand / neuromuscular (fine tremor, palm warmth/moisture, deep tendon reflex, proximal strength, CTS signs)", form["pe_hand_neuro"]),
        _line("Cardiovascular (rhythm regular/irregular, rate, murmur, JVP, edema)", form["pe_cv"]),
        _line("Skin / hair (warm-moist vs dry-coarse, pretibial myxedema, hair pattern, nails onycholysis)", form["pe_skin"]),
        _line("Other focused positives or important negatives", form["pe_other"]),
    ]))

    parts.append(_section("Prior data on file (free-text — physician's record of prior outside labs/images)", [
        _line("Prior thyroid labs (TSH/FT4/T3 values + date)", form["prior_labs"]),
        _line("Prior imaging (US/scan reports + date)", form["prior_imaging"]),
        _line("Prior FNA Bethesda category if any (date + result)", form["prior_fna"]),
    ]))

    has_screens = bool(screenshots)
    parts.append(
        f"\n## Lab/imaging screenshots\n"
        f"- Attached: {'yes — extract clearly readable values and incorporate them into Section 3 (Prior data on file)' if has_screens else 'no'}\n"
    )

    parts.append(
        "\n# Task\n"
        "Produce the 4-section thyroid initial-visit note following the OUTPUT FORMAT in the system prompt EXACTLY. "
        "Reminders:\n"
        "- Section 1 (Red flags): output ONLY if a red-flag pattern is triggered per the rules. Otherwise OMIT entirely.\n"
        "- Section 4A orders: phenotype-driven per the rules — do NOT output a blanket panel.\n"
        "- Blank / `(not documented / skipped)` fields → omit those sentences, do not pad.\n"
        "- Output is English only.\n"
    )

    text = "\n".join(parts)
    content: list[dict] = [{"type": "text", "text": text}]
    for path in screenshots:
        content.append(_image_block(path))
    if screenshots:
        content.append({
            "type": "text",
            "text": (
                "↑ Images above are screenshots of prior labs / imaging reports. "
                "Extract only clearly legible values. Incorporate into Section 3."
            ),
        })
    return content


# ──────────────────────────────────────────────────────────────────────────
# Generation
# ──────────────────────────────────────────────────────────────────────────


def _filename_for(chart_number: str | None) -> str:
    raw = (chart_number or "").strip()
    safe = re.sub(r"[^\w\-_.()一-鿿]+", "_", raw)[:60] if raw else ""
    ts = datetime.now(TAIPEI_TZ).strftime("%Y-%m-%d_%H_%M")
    return f"{safe}_thyroid_{ts}.md" if safe else f"thyroid_{ts}.md"


def generate(*values, screenshots, api_key_override, chart_number):
    api_key = (api_key_override or "").strip() or os.environ.get("ANTHROPIC_API_KEY", "").strip()
    if not api_key:
        yield ("❌ ANTHROPIC_API_KEY not set. Add it as a Space secret or paste it in the API key field.", "")
        return

    form = dict(zip(FORM_KEYS, values))

    paths: list[str] = []
    if screenshots:
        for f in screenshots:
            p = f if isinstance(f, str) else getattr(f, "name", None) or getattr(f, "path", None)
            if p:
                paths.append(p)

    try:
        client = anthropic.Anthropic(api_key=api_key)
        content = build_user_message(form, paths)
        msg = client.messages.create(
            model=MODEL,
            max_tokens=4096,
            system=[{
                "type": "text",
                "text": SYSTEM_PROMPT,
                "cache_control": {"type": "ephemeral"},
            }],
            messages=[{"role": "user", "content": content}],
        )
        out = "".join(block.text for block in msg.content if hasattr(block, "text"))
    except anthropic.APIError as e:
        yield (f"❌ Anthropic API error: {e}", "")
        return
    except Exception as e:  # noqa: BLE001
        yield (f"❌ Error: {e}", "")
        return

    if not gdrive.is_configured():
        yield (out, "_ℹ️ Google Drive sync not configured — set GDRIVE_* secrets to enable._")
        return

    filename = _filename_for(chart_number)
    yield (out, f"⏳ Uploading `{filename}` to Google Drive…")
    result = gdrive.upload_markdown(out, filename)
    if result.ok:
        link = f" — [open in Drive]({result.link})" if result.link else ""
        yield (out, f"✅ Saved to Google Drive as `{filename}`{link}")
    else:
        err = result.skipped_reason or result.error or "unknown error"
        yield (out, f"⚠️ Google Drive upload failed: {err}")


# ──────────────────────────────────────────────────────────────────────────
# Form schema (order must match build_user_message access)
# ──────────────────────────────────────────────────────────────────────────

FORM_KEYS = [
    # Visit context
    "age", "sex", "clinic_type", "visit_datetime", "chief_concern", "onset",
    # Phenotype tags
    "pt_hyper", "pt_hypo", "pt_nodule", "pt_abnormal_lab",
    # ROS
    "ros_hyper", "ros_hypo", "ros_neck", "ros_eye", "ros_pain",
    # PMH / FH
    "pmh_thyroid", "pmh_autoimmune", "pmh_cv", "pmh_other",
    "fh_thyroid", "fh_autoimmune",
    # Meds / exposure
    "meds_thyroid", "meds_other_thyroid", "exposure_iodine", "exposure_radiation",
    # Female-specific
    "female_pregnant", "female_menses",
    # PE
    "pe_vitals", "pe_general", "pe_eye", "pe_neck_inspect", "pe_neck_palp",
    "pe_neck_ausc", "pe_lymph", "pe_hand_neuro", "pe_cv", "pe_skin", "pe_other",
    # Prior data
    "prior_labs", "prior_imaging", "prior_fna",
]


# ──────────────────────────────────────────────────────────────────────────
# Download
# ──────────────────────────────────────────────────────────────────────────


def save_md(output_text: str, chart_number: str) -> str | None:
    if not output_text or output_text.startswith("❌"):
        return None
    filename = _filename_for(chart_number)
    path = Path(tempfile.gettempdir()) / filename
    path.write_text(output_text, encoding="utf-8")
    return str(path)


# ──────────────────────────────────────────────────────────────────────────
# Drafts (GDrive-synced)
# ──────────────────────────────────────────────────────────────────────────


def _file_to_blob(filepath: str) -> dict | None:
    p = Path(filepath)
    if not p.exists():
        return None
    mime, _ = mimetypes.guess_type(filepath)
    return {
        "name": p.name,
        "mime": mime or "image/png",
        "data_b64": base64.standard_b64encode(p.read_bytes()).decode("ascii"),
    }


def _blob_to_filepath(blob: dict) -> str | None:
    data_b64 = blob.get("data_b64") or ""
    if not data_b64:
        return None
    mime = blob.get("mime") or "image/png"
    suffix = mimetypes.guess_extension(mime) or ".png"
    fd, path = tempfile.mkstemp(prefix="thyroid_draft_", suffix=suffix)
    with os.fdopen(fd, "wb") as f:
        f.write(base64.standard_b64decode(data_b64))
    return path


def refresh_drafts_list():
    if not gdrive.is_configured():
        return gr.update(choices=[], value=None)
    return gr.update(choices=gdrive.list_drafts())


def save_draft_handler(chart_number, screenshots, personal_notes, *values):
    if not gdrive.is_configured():
        return gr.update(), "❌ GDrive sync not configured on this Space."
    if not chart_number or not chart_number.strip():
        return gr.update(), "⚠️ Enter a chart number first."

    name = chart_number.strip()
    paths = []
    if screenshots:
        for f in screenshots:
            p = f if isinstance(f, str) else getattr(f, "name", None) or getattr(f, "path", None)
            if p:
                paths.append(p)

    blobs = [b for b in (_file_to_blob(p) for p in paths) if b]

    payload = {
        "version": 1,
        "kind": "thyroid_initial_visit",
        "chart_number": name,
        "form": dict(zip(FORM_KEYS, values)),
        "personal_notes": personal_notes or "",
        "screenshots": blobs,
        "ts": datetime.utcnow().isoformat() + "Z",
    }
    result = gdrive.save_draft(name, payload)
    if not result.ok:
        return gr.update(), f"⚠️ Save failed: {result.error or result.skipped_reason}"

    names = gdrive.list_drafts()
    img_note = f" (with {len(blobs)} screenshot{'s' if len(blobs) != 1 else ''})" if blobs else ""
    return gr.update(choices=names, value=name), f"✅ Saved draft `{name}`{img_note} to Drive."


def load_draft_handler(name):
    n_form = len(FORM_KEYS)
    # outputs: chart_number + form (n_form) + personal_notes + screenshots + status
    n_total = 1 + n_form + 1 + 1 + 1
    noop = [gr.update() for _ in range(n_total - 1)]

    if not name:
        return noop + ["⚠️ Pick a draft first."]
    if not gdrive.is_configured():
        return noop + ["❌ GDrive sync not configured."]

    payload = gdrive.load_draft(name)
    if not payload:
        return noop + [f"⚠️ Draft `{name}` not found in Drive."]

    chart_number = payload.get("chart_number") or name
    form = payload.get("form", {})
    form_values = [form.get(k, "") if not k.startswith("pt_") else bool(form.get(k, False)) for k in FORM_KEYS]
    personal_notes = payload.get("personal_notes") or ""

    screenshot_paths = []
    for blob in payload.get("screenshots") or []:
        path = _blob_to_filepath(blob)
        if path:
            screenshot_paths.append(path)

    n_imgs = len(screenshot_paths)
    img_note = f" (with {n_imgs} screenshot{'s' if n_imgs != 1 else ''})" if n_imgs else ""
    return (
        [chart_number]
        + form_values
        + [personal_notes]
        + [screenshot_paths or None]
        + [f"📥 Loaded draft `{name}`{img_note}."]
    )


def delete_draft_handler(name):
    if not name:
        return gr.update(), "⚠️ Pick a draft first."
    if not gdrive.is_configured():
        return gr.update(), "❌ GDrive sync not configured."
    result = gdrive.delete_draft(name)
    if not result.ok:
        return gr.update(), f"⚠️ Delete failed: {result.error}"
    return gr.update(choices=gdrive.list_drafts(), value=None), f"🗑️ Deleted draft `{name}`."


# ──────────────────────────────────────────────────────────────────────────
# ACR TI-RADS calculator (pure Python — no API call)
# ──────────────────────────────────────────────────────────────────────────

TIRADS_OPTIONS = {
    "composition": [
        ("Cystic or almost completely cystic", 0),
        ("Spongiform", 0),
        ("Mixed cystic and solid", 1),
        ("Solid or almost completely solid", 2),
    ],
    "echogenicity": [
        ("Anechoic", 0),
        ("Hyperechoic or isoechoic", 1),
        ("Hypoechoic", 2),
        ("Very hypoechoic", 3),
    ],
    "shape": [
        ("Wider-than-tall", 0),
        ("Taller-than-wide", 3),
    ],
    "margin": [
        ("Smooth", 0),
        ("Ill-defined", 0),
        ("Lobulated or irregular", 2),
        ("Extra-thyroidal extension", 3),
    ],
    "foci": [
        ("None or large comet-tail artifacts", 0),
        ("Macrocalcifications", 1),
        ("Peripheral (rim) calcifications", 2),
        ("Punctate echogenic foci", 3),
    ],
}


def _tirads_level(total: int) -> tuple[str, str]:
    if total == 0:
        return "TR1 (benign)", "No FNA. No follow-up needed."
    if total == 2:
        return "TR2 (not suspicious)", "No FNA. No follow-up needed."
    if total == 3:
        return "TR3 (mildly suspicious)", "FNA if ≥ 2.5 cm. Follow-up US if ≥ 1.5 cm."
    if total in (4, 5, 6):
        return "TR4 (moderately suspicious)", "FNA if ≥ 1.5 cm. Follow-up US if ≥ 1.0 cm."
    # total ≥ 7
    return "TR5 (highly suspicious)", "FNA if ≥ 1.0 cm. Follow-up US if ≥ 0.5 cm."


def tirads_compute(composition, echogenicity, shape, margin, foci_selected, size_cm):
    # foci_selected is multi-select list of labels
    pts = {
        "composition": dict(TIRADS_OPTIONS["composition"]).get(composition, 0),
        "echogenicity": dict(TIRADS_OPTIONS["echogenicity"]).get(echogenicity, 0),
        "shape": dict(TIRADS_OPTIONS["shape"]).get(shape, 0),
        "margin": dict(TIRADS_OPTIONS["margin"]).get(margin, 0),
    }
    foci_dict = dict(TIRADS_OPTIONS["foci"])
    foci_pts = sum(foci_dict.get(lbl, 0) for lbl in (foci_selected or []))
    pts["foci"] = foci_pts

    total = sum(pts.values())
    level, fna_rule = _tirads_level(total)

    # Size-based action
    try:
        size = float(size_cm) if size_cm else None
    except ValueError:
        size = None

    if size is None:
        action = "Nodule size not entered — enter size to compute FNA threshold."
    else:
        if total == 0 or total == 2:
            action = f"Size {size} cm: no FNA needed (benign / not suspicious category)."
        elif total == 3:
            if size >= 2.5:
                action = f"Size {size} cm ≥ 2.5 cm → **FNA recommended**."
            elif size >= 1.5:
                action = f"Size {size} cm ≥ 1.5 cm → **follow-up US at 1, 3, 5 years**, no FNA."
            else:
                action = f"Size {size} cm: no FNA, no routine follow-up."
        elif total in (4, 5, 6):
            if size >= 1.5:
                action = f"Size {size} cm ≥ 1.5 cm → **FNA recommended**."
            elif size >= 1.0:
                action = f"Size {size} cm ≥ 1.0 cm → **follow-up US at 1, 2, 3, 5 years**, no FNA."
            else:
                action = f"Size {size} cm: no FNA, no routine follow-up."
        else:  # ≥ 7 → TR5
            if size >= 1.0:
                action = f"Size {size} cm ≥ 1.0 cm → **FNA recommended**."
            elif size >= 0.5:
                action = f"Size {size} cm ≥ 0.5 cm → **annual follow-up US × 5 years**, no FNA."
            else:
                action = f"Size {size} cm < 0.5 cm: no FNA, no routine follow-up."

    breakdown = "\n".join([
        f"- Composition: {composition} → {pts['composition']} pt",
        f"- Echogenicity: {echogenicity} → {pts['echogenicity']} pt",
        f"- Shape: {shape} → {pts['shape']} pt",
        f"- Margin: {margin} → {pts['margin']} pt",
        f"- Echogenic foci: {', '.join(foci_selected) if foci_selected else 'none'} → {foci_pts} pt",
    ])
    summary = (
        f"## ACR TI-RADS Result\n\n"
        f"**Total points: {total}** → **{level}**\n\n"
        f"FNA rule for this category: {fna_rule}\n\n"
        f"**Action: {action}**\n\n"
        f"---\n\n"
        f"### Point breakdown\n{breakdown}\n"
    )
    return summary


# ──────────────────────────────────────────────────────────────────────────
# UI
# ──────────────────────────────────────────────────────────────────────────


def tb(label: str, lines: int = 1, placeholder: str = "") -> gr.Textbox:
    return gr.Textbox(label=label, lines=lines, placeholder=placeholder or "leave blank or 'nil' to skip")


def build_initial_visit_tab(inputs: dict) -> tuple[
    gr.components.Component,  # screenshots
    gr.components.Component,  # api_key_box
    gr.components.Component,  # draft_name_box
    gr.components.Component,  # personal_notes_box
    gr.components.Component,  # generate_btn
    gr.components.Component,  # output
    gr.components.Component,  # upload_status
    gr.components.Component,  # download_btn
    gr.components.Component,  # download_file
    gr.components.Component,  # draft_select
    gr.components.Component,  # save_draft_btn
    gr.components.Component,  # load_draft_btn
    gr.components.Component,  # delete_draft_btn
    gr.components.Component,  # refresh_btn
    gr.components.Component,  # draft_status
]:
    gr.Markdown(
        "# 🦋 Thyroid Initial-Visit Note Generator\n"
        "Endocrine / family-medicine outpatient · 4-section structured output · Not for unverified clinical use."
    )

    with gr.Accordion("⚙️ API key (optional override)", open=False):
        api_key_box = gr.Textbox(
            label="Anthropic API key",
            type="password",
            placeholder="leave blank to use Space secret ANTHROPIC_API_KEY",
        )

    with gr.Accordion("💾 Patient chart number + drafts (synced to Google Drive)", open=True):
        gr.Markdown(
            "_Chart number is the single identifier — draft name in Drive, filename of generated `.md`, "
            "and how you find your in-progress cases across devices. Save here, switch device, "
            "pick from dropdown, Load, continue._"
        )
        with gr.Row():
            draft_name_box = gr.Textbox(
                label="Patient chart number (病歷號)",
                placeholder="e.g. 1029332",
                scale=3,
            )
            save_draft_btn = gr.Button("💾 Save draft", scale=1, variant="secondary")
        with gr.Row():
            draft_select = gr.Dropdown(
                label="📋 Existing drafts (in Drive)",
                choices=[],
                interactive=True,
                scale=3,
            )
            refresh_btn = gr.Button("🔄 Refresh", scale=1, size="sm")
        with gr.Row():
            load_draft_btn = gr.Button("📥 Load selected", variant="secondary")
            delete_draft_btn = gr.Button("🗑️ Delete selected", variant="stop")
        draft_status = gr.Markdown("")

    with gr.Accordion("1. Visit context", open=True):
        with gr.Row():
            inputs["age"] = tb("Age", placeholder="e.g. 42")
            inputs["sex"] = gr.Dropdown(["female", "male", ""], label="Sex", value="")
            inputs["clinic_type"] = gr.Dropdown(
                ["endocrine OPD", "family medicine OPD", "internal medicine OPD", "general OPD", ""],
                label="Clinic type", value="",
            )
        inputs["visit_datetime"] = tb("Visit date/time", placeholder="e.g. 2026-06-01 14:30")
        inputs["chief_concern"] = tb("Chief concern", lines=2, placeholder="e.g. abnormal TSH on health check; neck mass; palpitations × 2 months")
        inputs["onset"] = tb("Symptom onset / time course", lines=2, placeholder="e.g. gradual over 3 months; first noted on health-check 2 weeks ago")

    with gr.Accordion("2. Phenotype tags (check what applies — drives Section 4 order selection)", open=True):
        gr.Markdown(
            "_Tick one or more. If you leave all unchecked, AI will infer from your notes "
            "and state the inferred phenotype at the top of Section 1._"
        )
        with gr.Row():
            inputs["pt_hyper"] = gr.Checkbox(label="Hyperthyroid symptoms", value=False)
            inputs["pt_hypo"] = gr.Checkbox(label="Hypothyroid symptoms", value=False)
            inputs["pt_nodule"] = gr.Checkbox(label="Goiter / thyroid nodule", value=False)
            inputs["pt_abnormal_lab"] = gr.Checkbox(label="Abnormal lab (health check)", value=False)

    with gr.Accordion("3. Review of systems (fill per phenotype — skip rest)", open=True):
        inputs["ros_hyper"] = tb(
            "Hyperthyroid axis",
            lines=3,
            placeholder="palpitation / tremor / heat intolerance / sweating / weight loss / appetite ↑ / diarrhea / insomnia / anxiety / proximal weakness / oligomenorrhea / chest pain / syncope",
        )
        inputs["ros_hypo"] = tb(
            "Hypothyroid axis",
            lines=3,
            placeholder="fatigue / cold intolerance / weight gain / constipation / dry skin / hair loss / facial-leg edema / hoarseness / cognitive slowing / depression / myalgia / menorrhagia / bradycardia symptoms",
        )
        inputs["ros_neck"] = tb(
            "Neck / compressive",
            lines=2,
            placeholder="neck mass / dysphagia / dyspnea / orthopnea / choking / hoarseness / rapid enlargement / Pemberton",
        )
        inputs["ros_eye"] = tb(
            "Eye (Graves)",
            lines=2,
            placeholder="proptosis / dry eye / diplopia / eye pain / visual decline / photophobia / lid lag / morning lid swelling",
        )
        inputs["ros_pain"] = tb(
            "Pain / inflammatory",
            lines=2,
            placeholder="thyroid tenderness / radiation to jaw-ear / recent URI / fever / postpartum 6 mo / transient hyper→hypo course",
        )

    with gr.Accordion("4. Past medical history + family history", open=False):
        inputs["pmh_thyroid"] = tb(
            "Prior thyroid disease (Graves / Hashimoto / nodule / cancer / thyroiditis / surgery / RAI)",
            lines=2,
        )
        inputs["pmh_autoimmune"] = tb("Other autoimmune disease")
        inputs["pmh_cv"] = tb("Cardiovascular / osteoporosis (AF / HF / CAD / HTN / fracture)")
        inputs["pmh_other"] = tb("Other relevant (psych / liver / renal / leukopenia / recent malignancy)")
        inputs["fh_thyroid"] = tb("Family h/o thyroid disease (Graves / Hashimoto / nodule / cancer / MEN2 / MTC)")
        inputs["fh_autoimmune"] = tb("Family h/o other autoimmune disease")

    with gr.Accordion("5. Medications, supplements, exposures (高權重 — 一定要問)", open=False):
        inputs["meds_thyroid"] = tb(
            "Thyroid drugs (levothyroxine dose+timing+co-meds; methimazole; PTU; T3 prep)",
            lines=2,
        )
        inputs["meds_other_thyroid"] = tb(
            "Other thyroid-relevant drugs (amiodarone / lithium / IFN-α / ICI / TKI / biotin / OCP / anti-seizure / dopamine / glucocorticoid)",
            lines=2,
        )
        inputs["exposure_iodine"] = tb(
            "Iodine exposure (recent CT-cath contrast / kelp / seaweed / kelp supplement / betadine / Chinese-medicine)",
        )
        inputs["exposure_radiation"] = tb(
            "Radiation exposure (childhood/young-adult head-neck radiotherapy)",
        )

    with gr.Accordion("6. Female-specific (skip if N/A)", open=False):
        inputs["female_pregnant"] = tb(
            "Pregnancy / planning / breastfeeding / postpartum (specify month)",
        )
        inputs["female_menses"] = tb("Menstrual / fertility status")

    with gr.Accordion("7. Physical examination", open=False):
        inputs["pe_vitals"] = tb(
            "Vital signs (BT°C / HR / RR / BP / SpO2 / weight kg)",
        )
        inputs["pe_general"] = tb("General appearance")
        inputs["pe_eye"] = tb(
            "Eye exam (lid retraction / lid lag / proptosis / EOM / visual acuity / conjunctival injection / periorbital edema)",
            lines=2,
        )
        inputs["pe_neck_inspect"] = tb(
            "Neck inspection (visible enlargement / asymmetry / scar / venous engorgement / moves with swallow)",
        )
        inputs["pe_neck_palp"] = tb(
            "Neck palpation (size / consistency / surface / tenderness / mobility / nodule loc-size-border / tracheal deviation / retrosternal extension)",
            lines=3,
        )
        inputs["pe_neck_ausc"] = tb("Neck auscultation (thyroid bruit)")
        inputs["pe_lymph"] = tb(
            "Cervical lymph nodes (location / size / consistency / mobility — flag hard, fixed, supraclavicular)",
            lines=2,
        )
        inputs["pe_hand_neuro"] = tb(
            "Hand / neuromuscular (fine tremor / palm warmth-moisture / DTR / proximal strength / CTS)",
            lines=2,
        )
        inputs["pe_cv"] = tb("Cardiovascular (rhythm / rate / murmur / JVP / edema)")
        inputs["pe_skin"] = tb(
            "Skin / hair (warm-moist vs dry-coarse / pretibial myxedema / hair pattern / onycholysis)",
        )
        inputs["pe_other"] = tb("Other focused positives / important negatives")

    with gr.Accordion("8. Prior data on file (free text)", open=False):
        inputs["prior_labs"] = tb(
            "Prior thyroid labs (TSH/FT4/T3 values + date)",
            lines=3,
            placeholder="e.g. 2026-05-15: TSH 0.02 (0.4-4.0), FT4 2.8 (0.8-1.8), T3 280 (60-180)",
        )
        inputs["prior_imaging"] = tb(
            "Prior imaging (US / scan reports + date)",
            lines=3,
        )
        inputs["prior_fna"] = tb(
            "Prior FNA Bethesda category (date + result)",
            lines=2,
        )
        screenshots = gr.File(
            label="📷 Prior lab / imaging report screenshots (upload from desktop)",
            file_count="multiple",
            file_types=["image"],
            type="filepath",
        )

    gr.Markdown(
        "### 📝 Personal notes (僅供你提醒自己用 — 不會送 AI、不會出現在病歷)\n"
        "_例如：跟學長確認 TRAb 是否要追、家屬聯絡時間、回診排程。會跟著草稿在 Drive 同步，"
        "但 Generate 時不會被當作病歷一部分。_"
    )
    personal_notes_box = gr.Textbox(
        label="",
        placeholder="個人提醒用，自由發揮",
        lines=4,
        show_label=False,
    )

    generate_btn = gr.Button("✍️ Generate thyroid initial-visit note", variant="primary", size="lg")
    output = gr.Textbox(label="Output (4-section structured note)", lines=30, show_copy_button=True)
    upload_status = gr.Markdown("")

    with gr.Row():
        download_btn = gr.Button("📄 Download as .md", variant="secondary")
        download_file = gr.File(label="Download link", interactive=False)

    return (
        screenshots, api_key_box, draft_name_box, personal_notes_box,
        generate_btn, output, upload_status, download_btn, download_file,
        draft_select, save_draft_btn, load_draft_btn, delete_draft_btn,
        refresh_btn, draft_status,
    )


def build_tirads_tab():
    gr.Markdown(
        "# 🔬 ACR TI-RADS Calculator\n"
        "Pure local calculation — does not call the API. Reference: ACR TI-RADS, "
        "Tessler et al. JACR 2017."
    )
    with gr.Row():
        composition = gr.Dropdown(
            choices=[lbl for lbl, _ in TIRADS_OPTIONS["composition"]],
            label="Composition (0-2 pt)",
            value="Solid or almost completely solid",
        )
        echogenicity = gr.Dropdown(
            choices=[lbl for lbl, _ in TIRADS_OPTIONS["echogenicity"]],
            label="Echogenicity (0-3 pt)",
            value="Hypoechoic",
        )
    with gr.Row():
        shape = gr.Dropdown(
            choices=[lbl for lbl, _ in TIRADS_OPTIONS["shape"]],
            label="Shape (0 or 3 pt)",
            value="Wider-than-tall",
        )
        margin = gr.Dropdown(
            choices=[lbl for lbl, _ in TIRADS_OPTIONS["margin"]],
            label="Margin (0-3 pt)",
            value="Smooth",
        )
    foci = gr.CheckboxGroup(
        choices=[lbl for lbl, _ in TIRADS_OPTIONS["foci"]],
        label="Echogenic foci (multi-select; points are additive — 0+1+2+3 pt)",
        value=[],
    )
    size = gr.Textbox(
        label="Nodule size (largest dimension, cm)",
        placeholder="e.g. 1.2",
    )
    compute_btn = gr.Button("Compute TI-RADS", variant="primary")
    result = gr.Markdown("")

    compute_btn.click(
        fn=tirads_compute,
        inputs=[composition, echogenicity, shape, margin, foci, size],
        outputs=result,
    )
    gr.Markdown(
        "---\n"
        "_Result is an aid for the physician. Final FNA decision and follow-up plan "
        "depend on clinical context (TSH suppression, patient comorbidity, prior US trend, etc.)._"
    )


# ──────────────────────────────────────────────────────────────────────────
# Assemble app
# ──────────────────────────────────────────────────────────────────────────


def build_app() -> gr.Blocks:
    with gr.Blocks(title="Thyroid Initial Visit", theme=gr.themes.Soft()) as demo:
        inputs: dict[str, gr.components.Component] = {}

        with gr.Tabs():
            with gr.Tab("📝 Initial visit note"):
                (
                    screenshots, api_key_box, draft_name_box, personal_notes_box,
                    generate_btn, output, upload_status, download_btn, download_file,
                    draft_select, save_draft_btn, load_draft_btn, delete_draft_btn,
                    refresh_btn, draft_status,
                ) = build_initial_visit_tab(inputs)

            with gr.Tab("🔬 ACR TI-RADS calculator"):
                build_tirads_tab()

        ordered_inputs = [inputs[k] for k in FORM_KEYS]

        def _generate_adapter(*args):
            yield from generate(
                *args[:-3],
                screenshots=args[-3],
                api_key_override=args[-2],
                chart_number=args[-1],
            )

        generate_btn.click(
            fn=_generate_adapter,
            inputs=ordered_inputs + [screenshots, api_key_box, draft_name_box],
            outputs=[output, upload_status],
        )

        download_btn.click(
            fn=save_md,
            inputs=[output, draft_name_box],
            outputs=download_file,
        )

        save_draft_btn.click(
            fn=save_draft_handler,
            inputs=[draft_name_box, screenshots, personal_notes_box] + ordered_inputs,
            outputs=[draft_select, draft_status],
        )
        load_draft_btn.click(
            fn=load_draft_handler,
            inputs=[draft_select],
            outputs=[draft_name_box] + ordered_inputs + [personal_notes_box, screenshots, draft_status],
        )
        delete_draft_btn.click(
            fn=delete_draft_handler,
            inputs=[draft_select],
            outputs=[draft_select, draft_status],
        )
        refresh_btn.click(
            fn=refresh_drafts_list,
            inputs=[],
            outputs=draft_select,
        )
        draft_select.change(
            fn=lambda v: v or gr.update(),
            inputs=[draft_select],
            outputs=[draft_name_box],
        )
        demo.load(fn=refresh_drafts_list, inputs=[], outputs=draft_select)

        gr.Markdown(
            "---\n"
            "_Template-driven output. Verify every line against the chart before pasting into the EMR. "
            "Drafts and generated notes are stored in your personal Google Drive (drive.file scope — "
            "the app only sees files it created)._"
        )

    return demo


# ──────────────────────────────────────────────────────────────────────────
# Launch
# ──────────────────────────────────────────────────────────────────────────


if __name__ == "__main__":
    app = build_app()
    launch_kwargs: dict = {}
    if int(gr.__version__.split(".")[0]) >= 5:
        launch_kwargs["ssr_mode"] = False
    if PASSWORD:
        launch_kwargs["auth"] = ("doctor", PASSWORD)
    app.launch(**launch_kwargs)