"""
GRACE Reader Study - expert radiologist annotation app (Gradio / Hugging Face Space).

Design goal: annotating one case requires ZERO scrolling and ZERO guessing.
Everything needed to judge and score a case is visible at once:
  - two-column layout, one row per rated item (left = reference CXR, right = the item to score);
  - every rating control sits directly below the image it refers to;
  - full scale legends printed inline, subjective definitions in bold red with a worked example;
  - per-image zoom / brightness / contrast controls (display-only, never touch stored data);
  - no best/worst ranking asked; rankings are DERIVED in the backend from per-item scores.

Auth: readers sign in the FIRST time with a one-time invite credential, then choose their own
username + password (stored hashed in the private dataset). After that they log in with their own
credentials; closing the tab does not sign them out.

Secrets/config come from environment variables (set them as HF Space secrets):
  HF_TOKEN            : HF token with write permission (for the private response/account dataset).
  RESPONSE_DATASET    : private dataset repo id (default DrSyedFaizan/grace-reader-responses).
  CASES_DATASET       : optional private dataset holding cases.json + images (snapshot at boot).
  READER_CREDENTIALS  : JSON {"invite_name": "invite_password", ...} used ONLY for first-time login.
  APP_SECRET          : secret string used to sign resume tokens (defaults derived from HF_TOKEN).

NEVER commit Keys.txt or any token into the Space repo (see .gitignore).
"""

import os
import io
import json
import time
import base64
import random
import hashlib
import hmac
import secrets
from pathlib import Path

import gradio as gr
from PIL import Image

# ----------------------------------------------------------------------------- config
APP_DIR = Path(__file__).parent
DATA_DIR = APP_DIR / "data"
LOCAL_BACKUP_DIR = APP_DIR / "local_responses"
LOCAL_BACKUP_DIR.mkdir(exist_ok=True)

HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()
RESPONSE_DATASET = os.environ.get("RESPONSE_DATASET", "DrSyedFaizan/grace-reader-responses").strip()
CASES_DATASET = os.environ.get("CASES_DATASET", "").strip()
APP_SECRET = os.environ.get("APP_SECRET", "").strip() or ("grace-" + hashlib.sha256(HF_TOKEN.encode()).hexdigest()[:16] if HF_TOKEN else "grace-dev-secret")

MAX_ITEMS = 6            # max anonymized items rated per case; keep cases small (2-3) to preserve no-scroll
DISPLAY_MAX_W = 820      # px, display-copy width cap (does not affect stored data)
SCHEMA_VERSION = 1
ACCOUNTS_PATH = "accounts/accounts.json"

try:
    _invite = json.loads(os.environ.get("READER_CREDENTIALS", "").strip() or "{}")
except Exception:
    _invite = {}
if not _invite:
    _invite = {"reader1": "changeme", "reader2": "changeme"}
    print("[WARN] READER_CREDENTIALS not set - using demo invites. Set the secret before the real study.")
INVITES = {str(k): str(v) for k, v in _invite.items()}     # one-time first-login credentials

# ----------------------------------------------------------------------------- HF api
try:
    from huggingface_hub import HfApi
    _api = HfApi(token=HF_TOKEN) if HF_TOKEN else None
except Exception as e:
    _api = None
    print(f"[WARN] huggingface_hub unavailable: {e}")


def _ensure_response_dataset():
    if not _api:
        return
    try:
        _api.create_repo(RESPONSE_DATASET, repo_type="dataset", private=True, exist_ok=True)
    except Exception as e:
        print(f"[WARN] could not ensure response dataset: {e}")


# ----------------------------------------------------------------------------- cases
def _maybe_pull_cases():
    if not CASES_DATASET or not _api:
        return
    try:
        from huggingface_hub import snapshot_download
        snapshot_download(CASES_DATASET, repo_type="dataset", local_dir=str(DATA_DIR),
                          token=HF_TOKEN, local_dir_use_symlinks=False)
        print(f"[info] pulled cases from {CASES_DATASET}")
    except Exception as e:
        print(f"[WARN] could not pull CASES_DATASET: {e}")


def load_cases():
    _maybe_pull_cases()
    cj = DATA_DIR / "cases.json"
    if not cj.exists():
        print("[WARN] no data/cases.json - run build_cases_example.py to generate a demo set.")
        return []
    with open(cj, "r", encoding="utf-8") as f:
        data = json.load(f)
    return data.get("cases", [])


CASES = load_cases()
N_CASES = len(CASES)
_IMG_CACHE = {}


def img_data_uri(rel_path):
    if not rel_path:
        return ""
    if rel_path in _IMG_CACHE:
        return _IMG_CACHE[rel_path]
    p = DATA_DIR / rel_path
    try:
        im = Image.open(p).convert("RGB")
        if im.width > DISPLAY_MAX_W:
            h = int(im.height * DISPLAY_MAX_W / im.width)
            im = im.resize((DISPLAY_MAX_W, h))
        buf = io.BytesIO()
        im.save(buf, format="PNG")
        uri = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
    except Exception as e:
        print(f"[WARN] image load failed {p}: {e}")
        uri = ""
    _IMG_CACHE[rel_path] = uri
    return uri


# ----------------------------------------------------------------------------- tokens / passwords / accounts
def make_token(name):
    return hmac.new(APP_SECRET.encode(), name.encode(), hashlib.sha256).hexdigest()


def valid_token(name, tok):
    return bool(name) and bool(tok) and hmac.compare_digest(make_token(name), tok)


def hash_pw(pw, salt):
    return hashlib.pbkdf2_hmac("sha256", pw.encode(), bytes.fromhex(salt), 100_000).hex()


def verify_pw(pw, rec):
    try:
        return hmac.compare_digest(hash_pw(pw, rec["salt"]), rec["hash"])
    except Exception:
        return False


def load_accounts():
    """Persistent account store: {"accounts": {user: {salt,hash,...}}, "claimed_invites": [...]}."""
    default = {"accounts": {}, "claimed_invites": []}
    if _api:
        try:
            from huggingface_hub import hf_hub_download
            fp = hf_hub_download(RESPONSE_DATASET, ACCOUNTS_PATH, repo_type="dataset", token=HF_TOKEN)
            with open(fp, "r", encoding="utf-8") as f:
                d = json.load(f)
            d.setdefault("accounts", {}); d.setdefault("claimed_invites", [])
            return d
        except Exception:
            pass
    lb = LOCAL_BACKUP_DIR / "accounts.json"
    if lb.exists():
        with open(lb, "r", encoding="utf-8") as f:
            return json.load(f)
    return default


def save_accounts(acc):
    payload = json.dumps(acc, ensure_ascii=False, indent=2)
    with open(LOCAL_BACKUP_DIR / "accounts.json", "w", encoding="utf-8") as f:
        f.write(payload)
    if _api:
        try:
            _api.upload_file(path_or_fileobj=payload.encode("utf-8"), path_in_repo=ACCOUNTS_PATH,
                             repo_id=RESPONSE_DATASET, repo_type="dataset",
                             commit_message="update accounts")
        except Exception as e:
            print(f"[WARN] account upload failed (kept local): {e}")


# ----------------------------------------------------------------------------- storage
def _resp_path(annotator):
    return f"responses/{annotator}.jsonl"


def load_existing_responses(annotator):
    records = []
    if _api:
        try:
            from huggingface_hub import hf_hub_download
            fp = hf_hub_download(RESPONSE_DATASET, _resp_path(annotator), repo_type="dataset", token=HF_TOKEN)
            with open(fp, "r", encoding="utf-8") as f:
                return [json.loads(l) for l in f if l.strip()]
        except Exception:
            pass
    lb = LOCAL_BACKUP_DIR / f"{annotator}.jsonl"
    if lb.exists():
        with open(lb, "r", encoding="utf-8") as f:
            records = [json.loads(l) for l in f if l.strip()]
    return records


def completed_case_ids(records):
    return {r["case_id"] for r in records if r.get("item_id") == "__case__"}


def save_records(annotator, new_records):
    existing = load_existing_responses(annotator)
    existing.extend(new_records)
    payload = "\n".join(json.dumps(r, ensure_ascii=False) for r in existing) + "\n"
    with open(LOCAL_BACKUP_DIR / f"{annotator}.jsonl", "w", encoding="utf-8") as f:
        f.write(payload)
    if _api:
        try:
            _api.upload_file(path_or_fileobj=payload.encode("utf-8"), path_in_repo=_resp_path(annotator),
                             repo_id=RESPONSE_DATASET, repo_type="dataset",
                             commit_message=f"responses {annotator} {int(time.time())}")
        except Exception as e:
            print(f"[WARN] HF upload failed (kept local backup): {e}")
    return existing


# ----------------------------------------------------------------------------- ordering
def item_order(annotator, case):
    items = list(case.get("items", []))
    seed = int(hashlib.sha256(f"{annotator}|{case['case_id']}".encode()).hexdigest(), 16) % (2**32)
    random.Random(seed).shuffle(items)
    return items


def first_unfinished(records):
    done = completed_case_ids(records)
    for i, c in enumerate(CASES):
        if c["case_id"] not in done:
            return i
    return N_CASES


# ----------------------------------------------------------------------------- UI text
def image_html(dom_id, uri):
    if not uri:
        return "<div class='imgcell'><div class='imgbox empty'>[no image]</div></div>"
    return f"""
<div class='imgcell'>
  <div class='imgbox'><img id='{dom_id}' src='{uri}' data-bright='1' data-contrast='1' data-zoom='1'></div>
  <div class='strip'>
    <button type='button' onclick="graceZoom('{dom_id}',-0.25)" title='zoom out'>&minus;</button>
    <span class='lbl'>zoom</span>
    <button type='button' onclick="graceZoom('{dom_id}',0.25)" title='zoom in'>+</button>
    <label class='lbl'>bright<input type='range' min='0.3' max='2.5' step='0.05' value='1'
        oninput="graceSet('{dom_id}','bright',this.value)"></label>
    <label class='lbl'>contrast<input type='range' min='0.3' max='2.5' step='0.05' value='1'
        oninput="graceSet('{dom_id}','contrast',this.value)"></label>
  </div>
</div>"""


APPROP_LEGEND = ("**Decision appropriateness (1-5):**  "
                 "1 = clearly inappropriate,  2 = probably inappropriate,  3 = borderline / unsure,  "
                 "4 = probably appropriate,  5 = clearly appropriate.")
GROUND_DEF = ("<span class='reddef'><b>Grounding relevance</b> = does the highlighted region sit on the "
              "actual finding? <b>Example:</b> a right-lower-lobe opacity with the highlight on the right "
              "lower lobe = <b>Relevant</b>; the same case with the highlight on the heart border = "
              "<b>Not relevant</b>.</span>")
APPROP_DEF = ("<span class='reddef'><b>Appropriate</b> = the system's choice to answer vs defer was the safe, "
              "correct call for this image. <b>Example:</b> a subtle, ambiguous nodule where the system "
              "<b>defers to a radiologist</b> = appropriate; an obvious large opacity it needlessly defers = "
              "inappropriate.</span>")

INTRO_DEFAULT = (
    "The item(s) below come from anonymized automated reading systems, shown in random order with all "
    "identifying names hidden. For each item you see the system's <b>answer</b> to the clinical question, "
    "whether it chose to <b>answer</b> or <b>defer to a radiologist</b>, and a <b>highlight</b> of the image "
    "region it relied on. The final row shows the <b>reference region</b> (ground-truth annotation) for "
    "comparison. You cannot tell which system produced which item, and that is intentional.")

# ---- on-screen instructions (the app is fully self-explanatory; no external doc needed) ----
LOGIN_OVERVIEW = (
    "### What this study is\n"
    "You are helping validate an AI system that reads chest X-rays. For each case the system answers a clinical "
    "question and either commits to its answer or **defers to a radiologist**. Your job is to judge each "
    "anonymized system on three things: is the **answer correct**, was its decision to **answer vs defer "
    "appropriate**, and did it **look at the right region**.\n\n"
    "### What you will do\n"
    "- Each case shows one **reference chest X-ray** (left) and one or more **anonymized system outputs** (right), "
    "in random order with all system names hidden.\n"
    "- Score each output on the three questions, answer one question about the reference annotation, then press "
    "**Save & Next**.\n"
    "- A **progress counter** shows how many of the fixed number of cases you have finished.\n"
    "- You can **stop at any time**; your work saves as you go and you resume exactly where you left off.\n"
    "- Full step-by-step instructions stay on screen the whole time (in the **How to complete each case** panel).\n\n"
    "### Signing in\n"
    "**First time:** sign in with the one-time invite name and password you were given, then choose your own "
    "username and password. **Returning:** use the username and password you created. Closing the tab does "
    "**not** sign you out.")

SETUP_NOTE = (
    "This is a **one-time setup**. Choose a username you will remember (it is how your work is saved). Your "
    "progress saves automatically as you go, and you can stop and resume anytime. If a username is taken, pick "
    "another.")

HOWTO = (
    "**Layout.** Left column = the reference chest X-ray, always shown for comparison. Right column = one "
    "anonymized system's output: its **decision** (answered or deferred), its **answer**, and a **highlight** of "
    "the region it used. The final row shows the **reference region (ground truth)**. System order is random and "
    "names are hidden on purpose.\n\n"
    "**For each item on the right, score three things:**\n"
    "1. **Answer correctness** - Correct / Incorrect / Indeterminate (cannot tell from this image).\n"
    "2. **Decision appropriateness (1-5)** - was answering vs deferring the safe, correct call? 1 = clearly "
    "inappropriate ... 5 = clearly appropriate. *Example:* deferring a subtle, ambiguous nodule is appropriate; "
    "needlessly deferring an obvious large opacity is not.\n"
    "3. **Grounding relevance** - did the highlight sit on the actual finding? Relevant / Partial / Not relevant. "
    "*Example:* highlight on a right-lower-lobe opacity = Relevant; highlight on the heart border for that case = "
    "Not relevant.\n"
    "The **note** box (optional) is for anything unusual.\n\n"
    "**One case-level question:** is the reference (ground-truth) annotation acceptable - yes / partial / no.\n\n"
    "**Viewing each image.** Under every image are **-/+** zoom buttons and **bright** / **contrast** sliders. "
    "These change only how you view that image; they never change your scores or any stored data.\n\n"
    "**Saving & resuming.** Press **Save & Next** to store this case and move on (all three scores per item and "
    "the case-level question are required). You can stop anytime and log back in to resume where you left off. "
    "Use **Logout** to end your session. You can collapse this panel while scoring and reopen it anytime.")

DONE_MSG = (
    "### All cases complete. Thank you.\n\n"
    "Your responses have been saved securely. You may close this tab now. If you are asked to review additional "
    "cases later, simply log back in with your username and password and you will continue from the new cases.")

# localStorage helpers (client-side)
SET_LS_JS = "(u,t)=>{ if(t){ localStorage.setItem('grace_reader_user',u); localStorage.setItem('grace_reader_token',t);} }"
GET_LS_JS = "()=>[localStorage.getItem('grace_reader_user')||'', localStorage.getItem('grace_reader_token')||'']"
CLR_LS_JS = "()=>{ localStorage.removeItem('grace_reader_user'); localStorage.removeItem('grace_reader_token'); location.reload(); }"

# ----------------------------------------------------------------------------- head (JS + CSS)
HEAD = """
<style>
  .imgcell { display:flex; flex-direction:column; gap:4px; }
  .imgbox { overflow:auto; max-height:270px; border:1px solid #d0d5dd; border-radius:6px; background:#0b0b0b; }
  .imgbox.empty { display:flex; align-items:center; justify-content:center; color:#999; height:120px; background:#f3f4f6; }
  .imgbox img { display:block; max-width:100%; }
  .strip { display:flex; align-items:center; gap:8px; flex-wrap:wrap; font-size:12px; }
  .strip button { width:26px; height:24px; font-weight:700; cursor:pointer; }
  .strip .lbl { color:#475467; }
  .strip input[type=range] { width:90px; vertical-align:middle; }
  .reddef { color:#c0261c; font-size:13px; display:block; margin:2px 0 6px; }
  .legend { font-size:13px; color:#344054; }
  .refcol { border-right:2px dashed #cbd5e1; padding-right:8px; }
  .colhead { font-weight:700; font-size:13px; color:#101828; margin-bottom:2px; }
  .answerbox { background:#eef2ff; border:1px solid #c7d2fe; border-radius:6px; padding:6px 8px; font-size:13px; margin:4px 0; }
</style>
<script>
  window.graceApplyFilter = function(id){
    var img = document.getElementById(id); if(!img) return;
    var b = img.dataset.bright||1, c = img.dataset.contrast||1, z = img.dataset.zoom||1;
    img.style.filter = 'brightness('+b+') contrast('+c+')';
    img.style.transform = 'scale('+z+')'; img.style.transformOrigin = 'top left';
  };
  window.graceSet = function(id, kind, val){
    var img = document.getElementById(id); if(!img) return;
    img.dataset[kind] = val; window.graceApplyFilter(id);
  };
  window.graceZoom = function(id, delta){
    var img = document.getElementById(id); if(!img) return;
    var z = parseFloat(img.dataset.zoom||1) + delta;
    if(z<0.2) z=0.2; if(z>6) z=6; img.dataset.zoom=z; window.graceApplyFilter(id);
  };
  window.graceTooltips = function(){
    var map = {
      'Correct':'model answer matches the reference / your read',
      'Incorrect':'model answer does not match','Indeterminate':'cannot tell from this image',
      'Relevant':'highlight sits on the actual finding','Partial':'highlight partly overlaps the finding',
      'Not relevant':'highlight is on the wrong region',
      '1':'clearly inappropriate','2':'probably inappropriate','3':'borderline / unsure',
      '4':'probably appropriate','5':'clearly appropriate',
      'yes':'reference annotation is acceptable','partial':'reference partly acceptable','no':'reference not acceptable'
    };
    document.querySelectorAll('label').forEach(function(l){
      var t=(l.textContent||'').trim(); if(map[t]) l.title=map[t];
    });
  };
  setInterval(function(){ try{ window.graceTooltips(); }catch(e){} }, 1500);
</script>
"""

# ----------------------------------------------------------------------------- app
def build_app():
    with gr.Blocks(title="GRACE Reader Study", head=HEAD, theme=gr.themes.Soft()) as demo:
        annotator_state = gr.State("")
        index_state = gr.State(0)
        order_state = gr.State([])
        pending_invite_state = gr.State("")     # invite name being claimed during first-login setup

        # hidden helpers for localStorage resume
        user_ls = gr.Textbox(visible=False)
        tok_ls = gr.Textbox(visible=False)
        tok_out = gr.Textbox(visible=False)     # token written to localStorage after login/setup
        user_out = gr.Textbox(visible=False)    # username written to localStorage (may differ from typed)

        # ------------------------------------------------------------- LOGIN VIEW
        with gr.Column(visible=True) as login_col:
            gr.Markdown("## GRACE Reader Study")
            gr.Markdown(LOGIN_OVERVIEW)
            name_in = gr.Textbox(label="Username (or first-time invite name)", placeholder="username")
            pass_in = gr.Textbox(label="Password", type="password")
            login_btn = gr.Button("Sign in", variant="primary")
            login_msg = gr.Markdown("")

        # ------------------------------------------------------------- FIRST-LOGIN SETUP VIEW
        with gr.Column(visible=False) as setup_col:
            gr.Markdown("## Welcome - set up your account")
            gr.Markdown(SETUP_NOTE)
            su_user = gr.Textbox(label="Choose a username (min 3 chars)")
            su_pw = gr.Textbox(label="Choose a password (min 6 chars)", type="password")
            su_confirm = gr.Textbox(label="Confirm password", type="password")
            setup_btn = gr.Button("Create account & start", variant="primary")
            setup_msg = gr.Markdown("")

        # ------------------------------------------------------------- STUDY VIEW
        with gr.Column(visible=False) as study_col:
            with gr.Row():
                gr.Markdown("## GRACE Reader Study")
                progress_md = gr.Markdown("0 / 0")
                logout_btn = gr.Button("Logout", scale=0)
            with gr.Accordion("How to complete each case (click to expand or collapse)", open=True):
                gr.Markdown(HOWTO)
            intro_html = gr.HTML("")     # per-case provenance note (shown before the resume line)
            gr.Markdown("You can stop and resume anytime. Your answers save the moment you press "
                        "**Save & Next**.")
            question_md = gr.Markdown("")

            with gr.Row():
                gr.Markdown("<div class='colhead refcol'>REFERENCE image (for comparison)</div>")
                gr.Markdown("<div class='colhead'>ITEM to score</div>")

            gr.Markdown(f"<div class='legend'>{APPROP_LEGEND}</div>")
            gr.HTML(APPROP_DEF)
            gr.HTML(GROUND_DEF)

            item_rows, ref_htmls, item_htmls, ans_mds = [], [], [], []
            corr_rs, appr_rs, grnd_rs, note_tbs = [], [], [], []
            for i in range(MAX_ITEMS):
                with gr.Group(visible=False) as row:
                    with gr.Row():
                        with gr.Column(scale=1):
                            ref_htmls.append(gr.HTML("", elem_classes="refcol"))
                        with gr.Column(scale=1):
                            item_htmls.append(gr.HTML(""))
                            ans_mds.append(gr.Markdown(""))
                            corr_rs.append(gr.Radio(["Correct", "Incorrect", "Indeterminate"],
                                                    label="Answer correctness"))
                            appr_rs.append(gr.Radio(["1", "2", "3", "4", "5"],
                                                    label="Decision appropriateness (1-5)"))
                            grnd_rs.append(gr.Radio(["Relevant", "Partial", "Not relevant"],
                                                    label="Grounding relevance"))
                            note_tbs.append(gr.Textbox(label="Note (optional)", lines=1))
                item_rows.append(row)

            gr.Markdown("<div class='colhead'>Reference region (ground truth)</div>")
            with gr.Row():
                gt_ref_html = gr.HTML("", elem_classes="refcol")
                gt_html = gr.HTML("")
            case_acceptable = gr.Radio(["yes", "partial", "no"],
                                       label="Is the reference (ground-truth) annotation acceptable for this case?")
            status_md = gr.Markdown("")
            save_btn = gr.Button("Save & Next", variant="primary")
            done_md = gr.Markdown("", visible=False)

        # ---- ordered output list for render (must match render_case return order).
        CASE_OUTPUTS = ([intro_html, question_md, gt_ref_html, gt_html, case_acceptable,
                         progress_md, done_md]
                        + ref_htmls + item_htmls + ans_mds
                        + corr_rs + appr_rs + grnd_rs + note_tbs + item_rows)

        def render_case(annotator, idx, records=None):
            if records is None:
                records = load_existing_responses(annotator)
            done_n = len(completed_case_ids(records))
            progress = f"**{done_n} / {N_CASES}** cases completed"
            if idx >= N_CASES:
                ups = [gr.update(value=""), gr.update(value=""), gr.update(value=""),
                       gr.update(value=""), gr.update(value=None),
                       gr.update(value=progress),
                       gr.update(value=DONE_MSG, visible=True)]
                ups += [gr.update(value="") for _ in ref_htmls]
                ups += [gr.update(value="") for _ in item_htmls]
                ups += [gr.update(value="") for _ in ans_mds]
                ups += [gr.update(value=None) for _ in corr_rs]
                ups += [gr.update(value=None) for _ in appr_rs]
                ups += [gr.update(value=None) for _ in grnd_rs]
                ups += [gr.update(value="") for _ in note_tbs]
                ups += [gr.update(visible=False) for _ in item_rows]
                return ups, []

            case = CASES[idx]
            ref_uri = img_data_uri(case.get("reference_image", ""))
            items = item_order(annotator, case)
            order_ids = [it["item_id"] for it in items]
            intro = case.get("intro") or INTRO_DEFAULT
            q = "### " + case.get("question", "Assess the finding in this chest X-ray.")

            ref_ups, item_ups, ans_ups = [], [], []
            corr_ups, appr_ups, grnd_ups, note_ups, row_ups = [], [], [], [], []
            for i in range(MAX_ITEMS):
                if i < len(items):
                    it = items[i]
                    ref_ups.append(gr.update(value=image_html(f"ref_{idx}_{i}", ref_uri)))
                    item_ups.append(gr.update(value=image_html(f"item_{idx}_{i}", img_data_uri(it.get("image", "")))))
                    dec = it.get("decision", "answer")
                    ans_ups.append(gr.update(value=(f"<div class='answerbox'><b>System decision:</b> "
                                                    f"{'ANSWERED' if dec=='answer' else 'DEFERRED to radiologist'}"
                                                    f"<br><b>Answer:</b> {it.get('answer','(none)')}</div>")))
                    corr_ups.append(gr.update(value=None)); appr_ups.append(gr.update(value=None))
                    grnd_ups.append(gr.update(value=None)); note_ups.append(gr.update(value=""))
                    row_ups.append(gr.update(visible=True))
                else:
                    ref_ups.append(gr.update(value="")); item_ups.append(gr.update(value=""))
                    ans_ups.append(gr.update(value="")); corr_ups.append(gr.update(value=None))
                    appr_ups.append(gr.update(value=None)); grnd_ups.append(gr.update(value=None))
                    note_ups.append(gr.update(value="")); row_ups.append(gr.update(visible=False))

            head = [
                gr.update(value=f"<p>{intro}</p>"),
                gr.update(value=q),
                gr.update(value=image_html(f"gtref_{idx}", ref_uri)),
                gr.update(value=image_html(f"gt_{idx}", img_data_uri(case.get("groundtruth_image", "")))),
                gr.update(value=None),
                gr.update(value=progress),
                gr.update(value="", visible=False),
            ]
            ups = head + ref_ups + item_ups + ans_ups + corr_ups + appr_ups + grnd_ups + note_ups + row_ups
            return ups, order_ids

        # ---- navigation helpers (all return the SAME ALL_NAV-shaped list) --------------
        def _blank_case():
            return [gr.update() for _ in CASE_OUTPUTS]

        def nav_login(msg=""):
            return ([gr.update(value=""), gr.update(value=""),          # tok_out, user_out
                     gr.update(visible=True), gr.update(visible=False), gr.update(visible=False),
                     "", 0, [], "",                                     # states
                     gr.update(value=msg), gr.update(value="")]         # login_msg, setup_msg
                    + _blank_case())

        def nav_setup(pending_invite, msg=""):
            return ([gr.update(value=""), gr.update(value=""),
                     gr.update(visible=False), gr.update(visible=True), gr.update(visible=False),
                     "", 0, [], pending_invite,
                     gr.update(value=""), gr.update(value=msg)]
                    + _blank_case())

        def nav_study(username, token):
            records = load_existing_responses(username)
            idx = first_unfinished(records)
            ups, order_ids = render_case(username, idx, records)
            return ([gr.update(value=token), gr.update(value=username),
                     gr.update(visible=False), gr.update(visible=False), gr.update(visible=True),
                     username, idx, order_ids, "",
                     gr.update(value=""), gr.update(value="")]
                    + ups)

        ALL_NAV = ([tok_out, user_out, login_col, setup_col, study_col,
                    annotator_state, index_state, order_state, pending_invite_state,
                    login_msg, setup_msg] + CASE_OUTPUTS)

        # ---- auth handlers -------------------------------------------------------------
        def do_login(name, pw):
            name = (name or "").strip()
            acc = load_accounts()
            custom = acc.get("accounts", {})
            claimed = set(acc.get("claimed_invites", []))
            if name in custom and verify_pw(pw, custom[name]):
                return nav_study(name, make_token(name))
            if name in INVITES and name not in claimed and pw == INVITES[name]:
                return nav_setup(name)          # first-time login -> account setup
            if name in INVITES and name in claimed:
                return nav_login("That first-time invite has already been used. Sign in with the username "
                                 "and password you created.")
            return nav_login("Invalid username or password.")

        def do_setup(new_user, new_pw, confirm, pending_invite):
            new_user = (new_user or "").strip()
            if not pending_invite:
                return nav_login("Session expired. Please sign in again.")
            if len(new_user) < 3:
                return nav_setup(pending_invite, "Choose a username of at least 3 characters.")
            if len(new_pw or "") < 6:
                return nav_setup(pending_invite, "Choose a password of at least 6 characters.")
            if new_pw != confirm:
                return nav_setup(pending_invite, "Passwords do not match.")
            acc = load_accounts()
            custom = acc.setdefault("accounts", {})
            claimed = acc.setdefault("claimed_invites", [])
            if new_user in custom or new_user in INVITES:
                return nav_setup(pending_invite, "That username is taken. Choose another.")
            salt = secrets.token_hex(8)
            custom[new_user] = {"salt": salt, "hash": hash_pw(new_pw, salt),
                                "invite": pending_invite, "created_ts": int(time.time())}
            if pending_invite not in claimed:
                claimed.append(pending_invite)
            save_accounts(acc)
            return nav_study(new_user, make_token(new_user))

        def do_auto_login(name, tok):
            name = (name or "").strip()
            if not valid_token(name, tok):
                return nav_login("")
            if name in load_accounts().get("accounts", {}):
                return nav_study(name, tok)
            return nav_login("")

        login_btn.click(do_login, [name_in, pass_in], ALL_NAV).then(
            None, [user_out, tok_out], None, js=SET_LS_JS)
        setup_btn.click(do_setup, [su_user, su_pw, su_confirm, pending_invite_state], ALL_NAV).then(
            None, [user_out, tok_out], None, js=SET_LS_JS)

        # ---- save & next ---------------------------------------------------------------
        def do_save(annotator, idx, order_ids, case_ok, *rating_vals):
            M = MAX_ITEMS
            corr, appr, grnd, note = (rating_vals[0:M], rating_vals[M:2*M],
                                      rating_vals[2*M:3*M], rating_vals[3*M:4*M])
            if idx >= N_CASES:
                return [idx, order_ids, gr.update(value="Nothing to save.")] + _blank_case()
            n_items = len(order_ids)
            missing = []
            for i in range(n_items):
                if not corr[i]: missing.append(f"item {i+1}: correctness")
                if not appr[i]: missing.append(f"item {i+1}: appropriateness")
                if not grnd[i]: missing.append(f"item {i+1}: grounding")
            if not case_ok:
                missing.append("reference-acceptable question")
            if missing:
                msg = "Please complete before saving: " + "; ".join(missing[:6]) + ("..." if len(missing) > 6 else "")
                return [idx, order_ids, gr.update(value=msg)] + _blank_case()

            case = CASES[idx]
            ts = int(time.time())
            new_records = []
            for i in range(n_items):
                new_records.append({
                    "schema_version": SCHEMA_VERSION, "annotator": annotator,
                    "case_id": case["case_id"], "item_id": order_ids[i], "shown_position": i,
                    "dims": {"answer_correctness": corr[i], "decision_appropriateness": appr[i],
                             "grounding_relevance": grnd[i], "note": note[i] or ""},
                    "ts": ts})
            new_records.append({
                "schema_version": SCHEMA_VERSION, "annotator": annotator,
                "case_id": case["case_id"], "item_id": "__case__",
                "dims": {"reference_acceptable": case_ok}, "shown_order": order_ids, "ts": ts})
            records = save_records(annotator, new_records)
            nxt = first_unfinished(records)
            ups, new_order = render_case(annotator, nxt, records)
            return [nxt, new_order, gr.update(value="Saved.")] + ups

        SAVE_INPUTS = ([annotator_state, index_state, order_state, case_acceptable]
                       + corr_rs + appr_rs + grnd_rs + note_tbs)
        SAVE_OUTPUTS = [index_state, order_state, status_md] + CASE_OUTPUTS
        save_btn.click(do_save, SAVE_INPUTS, SAVE_OUTPUTS)

        # ---- logout --------------------------------------------------------------------
        def do_logout():
            return (gr.update(visible=True), gr.update(visible=False), gr.update(visible=False),
                    "", 0, [], "")
        logout_btn.click(do_logout, None,
                         [login_col, setup_col, study_col, annotator_state, index_state,
                          order_state, pending_invite_state]).then(None, None, None, js=CLR_LS_JS)

        # ---- boot / resume -------------------------------------------------------------
        demo.load(None, None, [user_ls, tok_ls], js=GET_LS_JS).then(
            do_auto_login, [user_ls, tok_ls], ALL_NAV)

    return demo


if __name__ == "__main__":
    _ensure_response_dataset()
    build_app().queue().launch(server_name="0.0.0.0", server_port=7860)