app.py

"""
app.py -- DECLARE user study: blind comparison of BASELINE / VANILLA-FT / DECLARE
on real spoken FairSpeech-style commands.

Single-file Flask app (matching the flat repo convention used by RailVaani),
consolidating what was previously split across server.py / models.py / nlu.py /
routing.py / logger.py into one file. Frontend lives in templates/index.html +
static/style.css + static/app.js.

Flow per trial:
  1. Participant is shown a command prompt to read aloud.
  2. They record themselves via the browser's microphone (MediaRecorder API).
  3. Audio is routed to ONE of three ASR models (Latin-square rotation, hidden
     from the participant).
  4. The transcript is passed through a rule-based intent/entity extractor.
  5. The participant judges intent-correctness and (if applicable)
     entity-correctness separately.
  6. Model identity is never revealed during the session.
  7. After all trials, a validated 4-item UMUX usability questionnaire is shown.

Run locally:   python app.py
Deploy on HF Spaces: see README.md (Docker SDK).

IMPORTANT -- session state is kept in an in-memory dict (SESSIONS), assuming a
SINGLE worker process. Do not scale to multiple gunicorn/Flask workers without
first moving session state to a shared store (Redis, a database, etc).
"""

import os
import re
import csv
import time
import uuid
import random
import subprocess
import tempfile
import threading

import pandas as pd
from flask import Flask, request, jsonify, render_template


# =============================================================================
# NLU -- rule-based intent + entity extraction
# =============================================================================
# IMPORTANT: This is a placeholder NLU layer, not a validated model. It exists
# so the study's task-outcome step ("did the assistant understand what you
# meant") has something concrete to show the participant. Before using results
# from this app in the actual paper, the intent/entity predictions shown to
# participants should be spot-checked, and ideally replaced with a proper
# trained classifier or manually verified gold labels.

INTENT_LABELS = [
    "COMMUNICATION & CALLING",
    "DEVICE CONTROL",
    "MUSIC & PLAYLIST CONTROL",
    "SOCIAL MEDIA OPERATIONS",
    "UNKNOWN",
]

INTENT_META = {
    "COMMUNICATION & CALLING": {
        "icon": "📞",
        "description": "Making calls, or sending texts/messages to a contact.",
    },
    "DEVICE CONTROL": {
        "icon": "⚙️",
        "description": "Adjusting phone settings — volume, camera, alarms, wifi, notifications.",
    },
    "MUSIC & PLAYLIST CONTROL": {
        "icon": "🎵",
        "description": "Playing, pausing, or managing songs and playlists.",
    },
    "SOCIAL MEDIA OPERATIONS": {
        "icon": "📱",
        "description": "Posting, sharing, or updating your social media status.",
    },
    "UNKNOWN": {
        "icon": "❓",
        "description": "Didn't clearly match any of the categories above.",
    },
}

_KEYWORDS = [
    ("COMMUNICATION & CALLING", [
        "call", "dial", "answer", "hang up", "decline",
        "text", "message", "whatsapp",
    ]),
    ("MUSIC & PLAYLIST CONTROL", [
        "play", "song", "music", "playlist", "album", "artist", "listen",
    ]),
    ("SOCIAL MEDIA OPERATIONS", [
        "share", "post", "status", "upload", "tag", "comment", "profile",
        "friend request",
    ]),
    ("DEVICE CONTROL", [
        "volume", "brightness", "wifi", "bluetooth", "alarm", "timer",
        "weather", "mute", "notification", "silence", "do not disturb",
        "unmute", "photo", "picture", "camera", "video", "record", "take a",
    ]),
]


def classify_intent(transcript: str) -> str:
    t = (transcript or "").lower()
    for label, kws in _KEYWORDS:
        if any(kw in t for kw in kws):
            return label
    return "UNKNOWN"


def guess_entity(transcript: str, intent: str) -> str:
    t = (transcript or "").lower().strip()
    if intent == "COMMUNICATION & CALLING":
        # Skip a leading possessive/article ("my", "the", etc.) before capturing
        # the name -- without this, "call my husband" incorrectly captured "my"
        # instead of "husband".
        m = re.search(r"(?:call|dial|answer)\w*\s+(?:my|your|his|her|our|their|the|a|an)?\s*([a-z]+)", t)
        if m:
            return m.group(1)
        m = re.search(r"(?:to|text)\s+(?:my|your|his|her|our|their|the|a|an)?\s*([a-z]+)(?:\s|$)", t)
        return m.group(1) if m else ""
    if intent == "MUSIC & PLAYLIST CONTROL":
        # Added rewind/pause/resume/stop/skip/repeat alongside "play" -- without
        # this, any non-"play" music command (e.g. "rewind the current music for
        # about ten seconds") returned no entity at all, even when the intent
        # itself was classified correctly. Truncation raised from 40->60 chars
        # so a trailing duration/detail phrase isn't accidentally cut off.
        m = re.search(r"(?:play|rewind|resume|repeat|replay|pause|skip)\s+(.+)", t)
        return m.group(1)[:60] if m else ""
    if intent == "DEVICE CONTROL":
        for kw in ["volume", "brightness", "wifi", "bluetooth", "alarm",
                    "timer", "weather", "camera", "photo", "video"]:
            if kw in t:
                return kw
    if intent == "SOCIAL MEDIA OPERATIONS":
        m = re.search(r"(?:share|post|upload)\s+(?:the\s+)?(.+?)(?:\s+to\s+|\s+on\s+|$)", t)
        return m.group(1)[:40] if m else ""
    return ""


def extract(transcript: str) -> dict:
    """Returns {'intent': str, 'icon': str, 'entity': str}"""
    intent = classify_intent(transcript)
    entity = guess_entity(transcript, intent)
    return {"intent": intent, "icon": INTENT_META[intent]["icon"], "entity": entity}


# =============================================================================
# Routing -- Latin-square blind assignment across the three ASR conditions
# =============================================================================
CONDITIONS = ["BASELINE", "VANILLA_FT", "DECLARE"]


def build_session_schedule(n_trials: int, seed: int = None) -> list:
    """Builds a schedule with exactly n_trials // 3 of each condition (3 each
    for the default 9-trial session), then fully shuffles the order -- a
    genuinely randomized sequence per session, not a repeating fixed rotation.
    Any leftover trials (if n_trials isn't a multiple of 3) are filled by
    randomly topping up from the condition list."""
    rng = random.Random(seed)
    base_count = n_trials // 3
    schedule = CONDITIONS * base_count
    remainder = n_trials - len(schedule)
    if remainder > 0:
        schedule += rng.sample(CONDITIONS, remainder)
    rng.shuffle(schedule)
    return schedule


def new_session_id() -> str:
    return uuid.uuid4().hex[:12]


# =============================================================================
# Logging -- trial results, participant demographics, UMUX responses
# =============================================================================
# NOTE: HF Spaces' local filesystem is ephemeral -- sync these CSVs to a
# private HF Dataset repo regularly if running more than a quick pilot.
LOG_PATH = os.environ.get("RESULTS_LOG_PATH", "results/trial_log.csv")
PARTICIPANTS_LOG_PATH = os.environ.get("PARTICIPANTS_LOG_PATH", "results/participants.csv")
UMUX_LOG_PATH = os.environ.get("UMUX_LOG_PATH", "results/umux_responses.csv")

# ── HF Dataset sync (optional) ────────────────────────────────────────────────
# Local CSVs remain the source of truth for reads (see log_* functions below);
# this additionally pushes each updated file to a private HF Dataset repo after
# every write, so results survive a Space restart/rebuild instead of living
# only in the container's ephemeral disk. Entirely optional: if HF_DATASET_REPO_ID
# or HF_SYNC_TOKEN aren't set, this is a silent no-op and the app behaves exactly
# as it did with local-CSV-only storage.
#
# Set these as Space secrets (Settings -> Variables and secrets):
#   HF_DATASET_REPO_ID = "your-username/declare-study-results"
#   HF_SYNC_TOKEN       = a write-scoped HF token
HF_DATASET_REPO_ID = os.environ.get("HF_DATASET_REPO_ID", "")
HF_SYNC_TOKEN = os.environ.get("HF_SYNC_TOKEN", "")
_hf_sync_enabled = bool(HF_DATASET_REPO_ID and HF_SYNC_TOKEN)

_hf_repo_ensured = False
_hf_repo_ensure_lock = threading.Lock()


def _ensure_hf_dataset_repo_exists():
    """Creates the dataset repo (private) if it doesn't already exist. Runs at
    most once per process, guarded by a lock since multiple upload threads
    could otherwise race to do this simultaneously."""
    global _hf_repo_ensured
    if _hf_repo_ensured:
        return
    with _hf_repo_ensure_lock:
        if _hf_repo_ensured:
            return
        try:
            from huggingface_hub import HfApi
            HfApi(token=HF_SYNC_TOKEN).create_repo(
                repo_id=HF_DATASET_REPO_ID, repo_type="dataset", exist_ok=True, private=True,
            )
        except Exception as e:
            print(f"[app.py] WARNING: could not ensure HF dataset repo exists: {e}")
        _hf_repo_ensured = True


def _sync_file_to_hf_dataset(local_path, repo_filename):
    """Uploads local_path to the configured HF dataset repo in a background
    thread, so this never adds latency to the participant-facing request that
    triggered it. Failures are logged, not raised -- a sync problem should
    never break the actual study flow, since the local CSV write already
    succeeded before this is called."""
    if not _hf_sync_enabled:
        return

    def _do_upload():
        try:
            _ensure_hf_dataset_repo_exists()
            from huggingface_hub import HfApi
            HfApi(token=HF_SYNC_TOKEN).upload_file(
                path_or_fileobj=local_path,
                path_in_repo=repo_filename,
                repo_id=HF_DATASET_REPO_ID,
                repo_type="dataset",
            )
            print(f"[app.py] Synced {repo_filename} -> HF dataset {HF_DATASET_REPO_ID}")
        except Exception as e:
            print(f"[app.py] WARNING: failed to sync {repo_filename} to HF dataset: {e}")

    threading.Thread(target=_do_upload, daemon=True).start()


print(f"[app.py] HF dataset sync: {'ENABLED -> ' + HF_DATASET_REPO_ID if _hf_sync_enabled else 'disabled (local CSV only)'}")

FIELDNAMES = [
    "timestamp", "session_id", "trial_index",
    "hash_name", "prompt_transcript", "prompt_domain",
    "condition",          # BASELINE / VANILLA_FT / DECLARE -- NOT shown to participant
    "asr_transcript", "predicted_intent", "predicted_entity",
    "trial_wer",           # word error rate, prompt_transcript vs asr_transcript
                            # (objective transcript-quality measure, to correlate
                            # against the participant's subjective correctness
                            # judgments below -- this is what RQ3 needs)
    "gold_intent", "gold_entity",   # pre-defined ground truth, NOT shown to participant
    "objective_intent_match",  # "yes" / "no" -- predicted_intent == gold_intent
    "objective_entity_match",  # "yes" / "no" / "not_applicable" -- machine-scored,
                                # independent of the participant's own judgment below
    "participant_intent_correct",  # "yes" / "no" -- the participant's own judgment
    "participant_entity_correct",  # "yes" / "no" / "not_applicable"
]

PARTICIPANT_FIELDNAMES = [
    "timestamp", "session_id", "gender", "age_group", "first_language",
    "recording_environment",  # Indoor / Outdoor / Not sure -- acoustic-condition
                               # covariate, NOT a demographic trait; captured
                               # because audio is discarded right after inference.
    "consent_given",
]

UMUX_FIELDNAMES = ["timestamp", "session_id", "q1", "q2", "q3", "q4", "umux_score"]

# UMUX (Usability Metric for User Experience) -- Finstad, K. (2010). The System
# Usability Scale and non-native English speakers. Journal of Usability Studies,
# 5(4), 185-191. A validated 4-item short-form of the 10-item SUS, reported
# reliability alpha=.94 and correlation with SUS r=.96 in the original study
# (replications found somewhat lower but still substantial correlations, ~.74-.81).
UMUX_ITEMS = [
    "This system's capabilities meet my requirements.",           # positive
    "Using this system is a frustrating experience.",             # negative
    "This system is easy to use.",                                # positive
    "I have to spend too much time correcting things with this system.",  # negative
]
UMUX_POSITIVE = [True, False, True, False]


def _ensure_file(path, fieldnames):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    if not os.path.exists(path):
        with open(path, "w", newline="") as f:
            csv.DictWriter(f, fieldnames=fieldnames).writeheader()


def log_participant(session_id, gender, age_group, first_language, recording_environment, consent_given):
    _ensure_file(PARTICIPANTS_LOG_PATH, PARTICIPANT_FIELDNAMES)
    with open(PARTICIPANTS_LOG_PATH, "a", newline="") as f:
        csv.DictWriter(f, fieldnames=PARTICIPANT_FIELDNAMES).writerow({
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "session_id": session_id, "gender": gender, "age_group": age_group,
            "first_language": first_language or "",
            "recording_environment": recording_environment or "Not sure",
            "consent_given": consent_given,
        })
    _sync_file_to_hf_dataset(PARTICIPANTS_LOG_PATH, "participants.csv")


def compute_wer(reference: str, hypothesis: str):
    """Standard word-level WER (edit distance / reference word count), the
    same definition used throughout the paper's benchmark evaluation, just
    computed per-trial here rather than corpus-level. Implemented directly
    (no jiwer dependency) to avoid adding another package to an already
    fragile requirements stack. Returns None if the reference has zero words
    (WER is undefined in that case, not zero)."""
    ref_words = (reference or "").strip().split()
    hyp_words = (hypothesis or "").strip().split()
    n, m = len(ref_words), len(hyp_words)
    if n == 0:
        return None
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n + 1):
        dp[i][0] = i
    for j in range(m + 1):
        dp[0][j] = j
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if ref_words[i - 1] == hyp_words[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
    return round(dp[n][m] / n, 4)


def score_entity_match(gold_entity, predicted_entity) -> str:
    """Substring containment in either direction, not exact match -- entities
    extracted by the rule-based NLU rarely match a gold label character-for-
    character (e.g. gold 'vito' vs predicted 'vito on iheartradio' should
    both count as a match). Returns 'not_applicable' when no entity was
    expected for this command at all.

    Coerces inputs to str defensively: pandas turns empty CSV cells into NaN
    (a float), which broke this function once already when a gold_entity
    column had a mix of blank and filled values."""
    gold = str(gold_entity if pd.notna(gold_entity) else "").strip().lower()
    pred = str(predicted_entity if pd.notna(predicted_entity) else "").strip().lower()
    if not gold:
        return "not_applicable"
    if not pred:
        return "no"
    return "yes" if (gold in pred or pred in gold) else "no"


def log_trial(session_id, trial_index, hash_name, prompt_transcript, prompt_domain,
              condition, asr_transcript, predicted_intent, predicted_entity,
              gold_intent, gold_entity,
              participant_intent_correct, participant_entity_correct):
    objective_intent_match = "yes" if predicted_intent == gold_intent else "no"
    objective_entity_match = score_entity_match(gold_entity, predicted_entity)

    _ensure_file(LOG_PATH, FIELDNAMES)
    with open(LOG_PATH, "a", newline="") as f:
        csv.DictWriter(f, fieldnames=FIELDNAMES).writerow({
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "session_id": session_id, "trial_index": trial_index,
            "hash_name": hash_name, "prompt_transcript": prompt_transcript,
            "prompt_domain": prompt_domain, "condition": condition,
            "asr_transcript": asr_transcript, "predicted_intent": predicted_intent,
            "predicted_entity": predicted_entity,
            "trial_wer": compute_wer(prompt_transcript, asr_transcript),
            "gold_intent": gold_intent, "gold_entity": gold_entity,
            "objective_intent_match": objective_intent_match,
            "objective_entity_match": objective_entity_match,
            "participant_intent_correct": participant_intent_correct,
            "participant_entity_correct": participant_entity_correct,
        })
    _sync_file_to_hf_dataset(LOG_PATH, "trial_log.csv")


def compute_umux_score(responses: list) -> float:
    """Each item scaled 0-6 (from a 1-7 response): positive items contribute
    (response - 1), negative items contribute (7 - response). Sum over the
    4 items (max 24) is normalised to 0-100."""
    total = 0
    for r, positive in zip(responses, UMUX_POSITIVE):
        total += (r - 1) if positive else (7 - r)
    return (total / 24) * 100


def log_umux(session_id, responses):
    score = compute_umux_score(responses)
    _ensure_file(UMUX_LOG_PATH, UMUX_FIELDNAMES)
    row = {"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "session_id": session_id}
    for i, r in enumerate(responses):
        row[f"q{i+1}"] = r
    row["umux_score"] = score
    with open(UMUX_LOG_PATH, "a", newline="") as f:
        csv.DictWriter(f, fieldnames=UMUX_FIELDNAMES).writerow(row)
    _sync_file_to_hf_dataset(UMUX_LOG_PATH, "umux_responses.csv")
    return score


# =============================================================================
# ASR models -- BASELINE / VANILLA_FT / DECLARE, with a mock-mode fallback
# =============================================================================
# *** YOU MUST SUPPLY YOUR OWN CHECKPOINTS FOR REAL INFERENCE ***
# Falls back to a clearly-labelled mock transcription if no real checkpoint is
# loaded, so the rest of the pipeline (routing, NLU, logging, UI) stays fully
# testable without any model weights.

ASR_MODEL_ID = os.environ.get("BASELINE_MODEL_ID", "stt_en_conformer_ctc_small_ls")

CHECKPOINT_PATHS = {
    "BASELINE": None,  # loaded via from_pretrained(ASR_MODEL_ID) below -- no path needed
    "VANILLA_FT": os.environ.get("VANILLA_FT_CKPT", "checkpoints/vanilla_ft.nemo"),
    "DECLARE": os.environ.get("DECLARE_CKPT", "checkpoints/declare.nemo"),
}

_loaded_models = {}
_nemo_available = False
try:
    import nemo.collections.asr as nemo_asr  # noqa: F401
    _nemo_available = True
except Exception as e:
    # IMPORTANT: log the real reason, don't swallow it silently. A broken
    # numba-cuda shim (see NUMBA_DISABLE_CUDA in the Dockerfile) is one known
    # cause of this import failing with an unrelated-looking error; there may
    # be others. Without this print, the app would fall back to mock mode
    # with no visible explanation in the logs.
    import traceback
    print(f"[app.py] NeMo import failed -- running in MOCK mode. Reason: {e}")
    traceback.print_exc()
    _nemo_available = False


def _load_model(condition: str):
    if condition in _loaded_models:
        return _loaded_models[condition]
    if not _nemo_available:
        _loaded_models[condition] = None
        return None

    import nemo.collections.asr as nemo_asr
    try:
        if condition == "BASELINE":
            model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=ASR_MODEL_ID)
        else:
            ckpt = CHECKPOINT_PATHS[condition]
            if not ckpt or not os.path.exists(ckpt):
                print(f"[app.py] Checkpoint not found for {condition}: {ckpt}")
                _loaded_models[condition] = None
                return None
            model = nemo_asr.models.ASRModel.restore_from(ckpt)
        model.eval()
        _loaded_models[condition] = model
        return model
    except Exception as e:
        print(f"[app.py] Failed to load {condition}: {e}")
        _loaded_models[condition] = None
        return None


_MOCK_NOISE = [
    lambda s: s,
    lambda s: s.replace("the", "da").replace("to", "too"),
    lambda s: " ".join(w[::-1] if len(w) > 4 else w for w in s.split()[:3]) + " " + " ".join(s.split()[3:]),
]


def _mock_transcribe(reference_text: str, condition: str) -> str:
    random.seed(hash((reference_text, condition)) % (2**32))
    noiser = random.choice(_MOCK_NOISE)
    return f"[MOCK-{condition}] " + noiser(reference_text)


def transcribe(condition: str, audio_path: str, reference_text: str = "") -> str:
    """Runs ASR for the given blind condition. The file at `audio_path` is
    deleted immediately after this function returns (see `finally`), whether
    inference succeeds, fails, or falls back to mock mode -- this is a hard
    participant-facing commitment ("your audio is never stored"), enforced
    here rather than left to the caller."""
    model = _load_model(condition)
    try:
        if model is None:
            return _mock_transcribe(reference_text, condition)
        try:
            out = model.transcribe([audio_path], verbose=False)[0]
            return out.text if hasattr(out, "text") else out
        except Exception as e:
            print(f"[app.py] Inference failed for {condition}: {e}")
            return _mock_transcribe(reference_text, condition)
    finally:
        try:
            if audio_path and os.path.exists(audio_path):
                os.remove(audio_path)
        except Exception as e:
            print(f"[app.py] Warning: could not delete audio file {audio_path}: {e}")


def model_status_report() -> dict:
    return {
        "nemo_available": _nemo_available,
        "baseline_loaded": _load_model("BASELINE") is not None,
        "vanilla_ft_loaded": _load_model("VANILLA_FT") is not None,
        "declare_loaded": _load_model("DECLARE") is not None,
    }


# =============================================================================
# Flask app -- routes
# =============================================================================
from werkzeug.exceptions import HTTPException

app = Flask(__name__)


@app.errorhandler(Exception)
def handle_uncaught_exception(e):
    """Without this, Flask's default behavior on an uncaught exception is to
    return an HTML error page -- which breaks every fetch() call on the
    client (res.json() fails with 'Unexpected token <, "<!DOCTYPE"...' since
    it's trying to parse HTML as JSON). This ensures the client always gets a
    parseable JSON error instead, and the real exception is still printed to
    the server log for debugging.

    IMPORTANT: HTTPException (404, 405, etc.) is deliberately excluded here
    and passed through to Flask/Werkzeug's own default handling. A previous
    version of this handler caught EVERYTHING including normal 404s from
    static file serving, which meant a request for app.js that hit a routing
    404 got converted into a JSON body with a 500 status instead of Flask's
    normal 404 response -- so the browser received JSON where it expected
    JavaScript, breaking every static asset request that didn't resolve
    exactly right. Only genuinely unexpected exceptions should reach here."""
    if isinstance(e, HTTPException):
        return e
    import traceback
    print(f"[app.py] Uncaught exception: {e}")
    traceback.print_exc()
    return jsonify({"error": f"Internal server error: {e}"}), 500

SAMPLE_COMMANDS_PATH = os.environ.get("SAMPLE_COMMANDS_PATH", "data/sample_commands.csv")
N_TRIALS_PER_SESSION = int(os.environ.get("N_TRIALS_PER_SESSION", "9"))

_samples_df = pd.read_csv(SAMPLE_COMMANDS_PATH, keep_default_na=False)

# In-memory session store: session_id -> dict. See module docstring re: single-worker constraint.
SESSIONS = {}


def _draw_trial_prompts(n):
    """Every participant sees the SAME fixed set of commands (data/sample_commands.csv
    is now curated to contain exactly the commands used in the study, with
    pre-defined gold_intent/gold_entity labels) -- only the presentation order
    is randomized per session, not which commands are shown. This removes
    command-difficulty as a confound when comparing conditions across
    participants: everyone judges the identical stimulus set."""
    df = _samples_df.sample(frac=1).reset_index(drop=True)  # shuffle order only
    return df.head(n).to_dict("records")


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/intent-legend")
def intent_legend():
    return jsonify({"labels": INTENT_LABELS, "meta": INTENT_META})


@app.route("/api/umux/items")
def umux_items():
    return jsonify({"items": UMUX_ITEMS})


@app.route("/api/session/start", methods=["POST"])
def session_start():
    data = request.get_json(force=True) or {}
    gender = data.get("gender")
    age_group = data.get("age_group")
    first_language = data.get("first_language", "")
    environment = data.get("environment")
    consent = data.get("consent", False)

    if not consent:
        return jsonify({"error": "Consent is required."}), 400
    if not gender or not age_group:
        return jsonify({"error": "Gender and age group are required."}), 400

    session_id = new_session_id()
    log_participant(
        session_id=session_id, gender=gender, age_group=age_group,
        first_language=first_language, recording_environment=environment,
        consent_given=True,
    )

    prompts = _draw_trial_prompts(N_TRIALS_PER_SESSION)
    schedule = build_session_schedule(len(prompts), seed=hash(session_id) % (2**32))

    SESSIONS[session_id] = {"prompts": prompts, "schedule": schedule, "trial_index": 0}

    first_prompt = prompts[0]
    return jsonify({
        "session_id": session_id,
        "total_trials": len(prompts),
        "trial_index": 0,
        "prompt_text": first_prompt["transcription"],
    })


@app.route("/api/trial/submit", methods=["POST"])
def trial_submit():
    session_id = request.form.get("session_id")
    audio_file = request.files.get("audio")

    if session_id not in SESSIONS:
        return jsonify({"error": "Invalid or expired session. Please refresh and start again."}), 400
    if audio_file is None:
        return jsonify({"error": "No audio received."}), 400

    state = SESSIONS[session_id]
    idx = state["trial_index"]
    prompt_row = state["prompts"][idx]
    condition = state["schedule"][idx]  # hidden from participant, never sent to client

    if _nemo_available:
        with _model_load_lock:
            status = _model_load_status.get(condition, "not_started")
        if status in ("not_started", "loading"):
            return jsonify({
                "error": "The ASR models are still warming up on the server. "
                         "Please wait about a minute and try again.",
                "still_loading": True,
            }), 503

    raw_fd, raw_path = tempfile.mkstemp(suffix=".webm")
    wav_fd, wav_path = tempfile.mkstemp(suffix=".wav")
    os.close(raw_fd)
    os.close(wav_fd)
    audio_file.save(raw_path)

    try:
        subprocess.run(
            ["ffmpeg", "-y", "-i", raw_path, "-ar", "16000", "-ac", "1", wav_path],
            check=True, capture_output=True,
        )
    except subprocess.CalledProcessError as e:
        if os.path.exists(raw_path):
            os.remove(raw_path)
        if os.path.exists(wav_path):
            os.remove(wav_path)
        return jsonify({"error": f"Audio conversion failed: {e.stderr.decode(errors='ignore')[:300]}"}), 500

    t0 = time.time()
    asr_transcript = transcribe(
        condition=condition, audio_path=wav_path, reference_text=prompt_row["transcription"],
    )  # transcribe() deletes wav_path itself once inference finishes
    prediction = extract(asr_transcript)
    elapsed = time.time() - t0

    if os.path.exists(raw_path):
        os.remove(raw_path)  # transcribe() only owns wav_path, not the original upload

    entity_applicable = bool(prediction["entity"])

    state["last_asr_transcript"] = asr_transcript
    state["last_prediction"] = prediction
    state["last_condition"] = condition
    state["last_prompt_row"] = prompt_row
    state["last_entity_applicable"] = entity_applicable

    return jsonify({
        "intent": prediction["intent"],
        "icon": prediction["icon"],
        "entity": prediction["entity"] if entity_applicable else None,
        "entity_applicable": entity_applicable,
        "elapsed_seconds": round(elapsed, 2),
    })


@app.route("/api/trial/judge", methods=["POST"])
def trial_judge():
    data = request.get_json(force=True) or {}
    session_id = data.get("session_id")
    intent_answer = data.get("intent_correct")
    entity_answer = data.get("entity_correct")

    if session_id not in SESSIONS:
        return jsonify({"error": "Invalid or expired session."}), 400

    state = SESSIONS[session_id]
    entity_applicable = state.get("last_entity_applicable", False)

    if intent_answer is None:
        return jsonify({"error": "intent_correct is required."}), 400
    if entity_applicable and entity_answer is None:
        return jsonify({"error": "entity_correct is required."}), 400

    intent_correct = "yes" if intent_answer == "Yes" else "no"
    entity_correct = (
        "not_applicable" if not entity_applicable
        else ("yes" if entity_answer == "Yes" else "no")
    )

    idx = state["trial_index"]
    prompt_row = state["last_prompt_row"]

    log_trial(
        session_id=session_id, trial_index=idx,
        hash_name=prompt_row["hash_name"], prompt_transcript=prompt_row["transcription"],
        prompt_domain=prompt_row["domain"], condition=state["last_condition"],
        asr_transcript=state["last_asr_transcript"],
        predicted_intent=state["last_prediction"]["intent"],
        predicted_entity=state["last_prediction"]["entity"],
        gold_intent=prompt_row.get("gold_intent", ""),
        gold_entity=prompt_row.get("gold_entity", ""),
        participant_intent_correct=intent_correct,
        participant_entity_correct=entity_correct,
    )

    state["trial_index"] += 1
    n_total = len(state["prompts"])

    if state["trial_index"] >= n_total:
        return jsonify({"session_complete": True})

    next_row = state["prompts"][state["trial_index"]]
    return jsonify({
        "session_complete": False,
        "trial_index": state["trial_index"],
        "total_trials": n_total,
        "prompt_text": next_row["transcription"],
    })


@app.route("/api/umux/submit", methods=["POST"])
def umux_submit():
    data = request.get_json(force=True) or {}
    session_id = data.get("session_id")
    responses = data.get("responses")

    if not responses or len(responses) != 4 or any(r is None for r in responses):
        return jsonify({"error": "All 4 questions must be answered."}), 400

    log_umux(session_id=session_id, responses=[int(r) for r in responses])
    return jsonify({"ok": True})


# ── Results download -- token-protected, since these are real participant  ──
# demographics + transcripts. Set ADMIN_TOKEN as a Space secret (Settings ->
# Variables and secrets); without it set, this defaults to a placeholder that
# you should change immediately.
ADMIN_TOKEN = os.environ.get("ADMIN_TOKEN", "change-me")

_DOWNLOADABLE_FILES = {
    "trial_log.csv": LOG_PATH,
    "participants.csv": PARTICIPANTS_LOG_PATH,
    "umux_responses.csv": UMUX_LOG_PATH,
}


@app.route("/admin/download/<filename>")
def download_results(filename):
    from flask import send_file

    token = request.args.get("token")
    if token != ADMIN_TOKEN:
        return jsonify({"error": "Invalid or missing token."}), 403

    if filename not in _DOWNLOADABLE_FILES:
        return jsonify({"error": f"Unknown file. Available: {list(_DOWNLOADABLE_FILES.keys())}"}), 404

    path = _DOWNLOADABLE_FILES[filename]
    if not os.path.exists(path):
        return jsonify({"error": f"{filename} doesn't exist yet -- no data has been logged."}), 404

    return send_file(path, as_attachment=True, download_name=filename)


@app.route("/api/debug/model-status")
def debug_model_status():
    return jsonify(model_status_report())


# Tracks each condition's loading state so routes can respond instantly with a
# clean "still warming up" message instead of blocking (and risking a proxy
# timeout / HTML error page) while a model is still loading in the background.
_model_load_status = {c: "not_started" for c in CONDITIONS}  # not_started / loading / ready / failed
_model_load_lock = threading.Lock()


def _background_load_all_models():
    for cond in CONDITIONS:
        with _model_load_lock:
            _model_load_status[cond] = "loading"
        _load_model(cond)
        with _model_load_lock:
            _model_load_status[cond] = "ready" if _loaded_models.get(cond) is not None else "failed"
    print(f"[app.py] Background model loading complete: {_model_load_status}")


# Started immediately at import time, but Flask's own app.run() (further down,
# in the __main__ block) is NOT blocked by this -- the server starts accepting
# connections right away, while models finish loading in parallel.
print("[app.py] Starting background ASR model loading (server accepts requests immediately) ...")
threading.Thread(target=_background_load_all_models, daemon=True).start()


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), threaded=True)