server.py

#!/usr/bin/env python3
"""Revolab Benchmark Explorer — FastAPI backend.

Run from the project root:
    python explorer/server.py
    # or
    uvicorn explorer.server:app --reload --port 9999
"""

from __future__ import annotations

import io
import json
import logging
import os
import shutil
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

import numpy as np
from dotenv import load_dotenv

load_dotenv()
from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles

PROJECT_ROOT = Path(__file__).resolve().parent.parent
EXPLORER_DIR = Path(__file__).resolve().parent

# Import from local vendored utils (no parent repo dependency)
from utils import read_manifest, build_freq_map, compute_cer, compute_rare_wer, make_common_words, decode_audio

TOP_N_COMMON = 500  # top-N most-frequent words considered "common" for rare-WER

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

STATIC_DIR = Path(__file__).resolve().parent / "static"

# Results directories — check local data/ first (for standalone repo), then parent results/
LOCAL_DATA_DIR = EXPLORER_DIR / "data"

# Private-split manifests are never committed to this repo (so the repo/Space can be
# public without exposing raw private transcripts). They're downloaded at startup from
# a private HF dataset repo, gated behind HF_TOKEN.
PRIVATE_RESULTS_DATASET = "Revolab/asr-benchmark-explorer-private-results"


def _ensure_private_data() -> None:
    target = LOCAL_DATA_DIR / "private"
    if any(target.glob("*.jsonl")):
        return
    token = os.environ.get("HF_TOKEN")
    if not token:
        logger.warning("HF_TOKEN not set — private-split leaderboard will be unavailable")
        return
    try:
        from huggingface_hub import snapshot_download
        snapshot_path = snapshot_download(
            repo_id=PRIVATE_RESULTS_DATASET,
            repo_type="dataset",
            token=token,
            allow_patterns=["data/private/*"],
        )
        src = Path(snapshot_path) / "data" / "private"
        target.mkdir(parents=True, exist_ok=True)
        for f in src.glob("*"):
            shutil.copy2(f, target / f.name)
        logger.info("Downloaded private-split data: %d files", len(list(target.glob("*"))))
    except Exception as exc:
        logger.warning("Could not download private-split data: %s", exc)


_ensure_private_data()

PUBLIC_DIR  = LOCAL_DATA_DIR / "public"  if LOCAL_DATA_DIR.joinpath("public").exists()  else PROJECT_ROOT / "results" / "public"
PRIVATE_DIR = LOCAL_DATA_DIR / "private" if LOCAL_DATA_DIR.joinpath("private").exists() else PROJECT_ROOT / "results" / "private"
LEGACY_DIR  = LOCAL_DATA_DIR / "malay-benchmark" if LOCAL_DATA_DIR.joinpath("malay-benchmark").exists() else PROJECT_ROOT / "results" / "malay-benchmark"

# Model IDs to exclude from all APIs (retired / experimental runs)
EXCLUDED_MODEL_IDS: frozenset = frozenset({
    "gemini-3.1-pro-preview",
    "qwen3-asr-0.6b-malay-s12000",
})

# Category → group mapping
TELEPHONY_CATS    = frozenset({"telephony-production", "telephony"})
OPENSOURCE_CATS   = frozenset({"fleurs", "commonvoice"})
SHORT_INPUT_CATS  = frozenset({"short-inputs"})

def _cat_group(cat: str) -> str:
    if cat in TELEPHONY_CATS:   return "telephony"
    if cat in OPENSOURCE_CATS:  return "opensource"
    if cat in SHORT_INPUT_CATS: return "short_inputs"
    return "domains"

app = FastAPI(title="Revolab Benchmark Explorer", docs_url=None, redoc_url=None)

# ─── Data storage ──────────────────────────────────────────────────────────────

# Public-only: (split, ref, ref_idx) → {model_id: record}  — used by Explorer
_pub_sample_map: dict[tuple[str, str, int], dict[str, dict]] = defaultdict(dict)
_pub_all_samples: list[dict] = []

# All records flat list — used by leaderboard
_all_records: list[dict] = []

_model_ids:  list[str] = []
_categories: list[str] = []
_splits:     list[str] = []

NOISE_LEVELS = ["clean", "moderate", "noisy"]

_common_words: frozenset = frozenset()  # populated after manifests load

# ─── Load manifests ────────────────────────────────────────────────────────────

def _load_dir(directory: Path, benchmark: str) -> list[dict]:
    records = []
    if not directory.exists():
        return records
    for path in sorted(directory.glob("*.jsonl")):
        batch = read_manifest(path)
        # Assign ref_idx per-file so duplicate refs in the same file get distinct indices.
        # Records across files are row-aligned (same audio), so the same ref_idx in
        # different files maps to the same audio sample.
        file_ref_counter: dict[tuple, int] = defaultdict(int)
        for r in batch:
            r["_benchmark"] = benchmark
            base = (r.get("split", ""), r.get("reference", ""))
            r["_ref_idx"] = file_ref_counter[base]
            file_ref_counter[base] += 1
        records.extend(batch)
        if batch:
            logger.info("  [%s] %s — %d records", benchmark, path.name, len(batch))
    return records


def _load_manifests() -> None:
    global _model_ids, _categories, _splits

    # Prefer results/public; fall back to legacy dir if public is empty
    pub_dir = PUBLIC_DIR
    if not any(PUBLIC_DIR.glob("*.jsonl")):
        pub_dir = LEGACY_DIR
        logger.info("results/public/ empty — using legacy %s", LEGACY_DIR)

    pub_records  = [r for r in _load_dir(pub_dir, "public")  if r.get("model_id") not in EXCLUDED_MODEL_IDS]
    priv_records = [r for r in _load_dir(PRIVATE_DIR, "private") if r.get("model_id") not in EXCLUDED_MODEL_IDS]

    _all_records.extend(pub_records)
    _all_records.extend(priv_records)

    model_set: set[str] = set()
    cat_set:   set[str] = set()
    split_set: set[str] = set()

    # Build public sample map (Explorer)
    for r in pub_records:
        key = (r["split"], r["reference"], r.get("_ref_idx", 0))
        _pub_sample_map[key][r["model_id"]] = r
        model_set.add(r["model_id"])
        if r.get("category"):
            cat_set.add(r["category"])
        split_set.add(r["split"])

    # Also collect model ids from private
    for r in priv_records:
        model_set.add(r["model_id"])

    for (split, ref, ref_idx), models in _pub_sample_map.items():
        first = next(iter(models.values()))
        model_wers = {
            model_id: round(
                (r.get("word_substitutions", 0) + r.get("word_insertions", 0) + r.get("word_deletions", 0))
                / max(r.get("num_ref_words", 1), 1) * 100,
                2,
            )
            for model_id, r in models.items()
        }
        _pub_all_samples.append({
            "split":          split,
            "ref":            ref,
            "ref_idx":        ref_idx,
            "category":       first.get("category", ""),
            "best_wer":       round(min(model_wers.values()), 1),
            "model_wers":     model_wers,
            "audio_length_s": first.get("audio_length_s", 0),
            "n_models":       len(models),
        })

    _model_ids[:] = sorted(model_set)
    _categories[:] = sorted(cat_set)
    _splits[:] = sorted(split_set)

    logger.info(
        "Loaded %d public + %d private records | %d models | %d public samples",
        len(pub_records), len(priv_records), len(_model_ids), len(_pub_all_samples),
    )


_load_manifests()

# Build rare-word vocabulary from all loaded references
def _build_common_words() -> None:
    global _common_words
    all_refs = [
        r.get("reference_normalized_final") or r.get("reference_normalized_ours") or ""
        for r in _all_records
    ]
    freq_map = build_freq_map(all_refs)
    _common_words = make_common_words(freq_map, TOP_N_COMMON)
    logger.info(
        "Rare-WER vocab: top %d common words (corpus size %d unique)",
        TOP_N_COMMON, len(freq_map),
    )


_build_common_words()

# ─── Noise tags ────────────────────────────────────────────────────────────────

_noise_tags: dict[str, dict[str, dict]] = {}


def _load_noise_tags() -> None:
    # Load public tags (prefer public over legacy)
    for d in (PUBLIC_DIR, LEGACY_DIR):
        path = d / "noise_tags.json"
        if path.exists():
            with open(path, encoding="utf-8") as f:
                _noise_tags.update(json.load(f))
            logger.info("Noise tags (public): %d splits from %s", len(_noise_tags), d)
            break
    # Load private tags (merged on top — same split/ref keys are overwritten but private
    # refs are unique so in practice they just add new entries)
    priv_path = PRIVATE_DIR / "noise_tags.json"
    if priv_path.exists():
        with open(priv_path, encoding="utf-8") as f:
            priv_tags = json.load(f)
        for split, refs in priv_tags.items():
            _noise_tags.setdefault(split, {}).update(refs)
        total_priv = sum(len(v) for v in priv_tags.values())
        logger.info("Noise tags (private): %d samples merged", total_priv)


_load_noise_tags()

# ─── Leaderboard ──────────────────────────────────────────────────────────────


def _compute_stats(records: list[dict]) -> dict | None:
    if not records:
        return None
    total_err = total_ref = total_subs = total_ins = total_dels = 0
    total_char_err = total_char_ref = 0
    total_audio = total_time = 0.0
    per_wers: list[float] = []
    by_cat: dict[str, dict] = defaultdict(lambda: {"err": 0, "ref": 0, "n": 0, "wers": []})

    for r in records:
        s  = r.get("word_substitutions", 0)
        i  = r.get("word_insertions", 0)
        d  = r.get("word_deletions", 0)
        nr = r.get("num_ref_words", 0)
        err = s + i + d
        total_err  += err
        total_ref  += nr
        total_subs += s
        total_ins  += i
        total_dels += d
        total_audio += r.get("audio_length_s", 0)
        total_time  += r.get("transcription_time_s", 0)
        if nr > 0:
            per_wers.append(err / nr * 100)
        cat = r.get("category", "")
        if cat:
            by_cat[cat]["err"] += err
            by_cat[cat]["ref"] += nr
            by_cat[cat]["n"]   += 1
            if nr > 0:
                by_cat[cat]["wers"].append(err / nr * 100)

        # CER: character-level edit distance on normalised text
        ref_norm  = r.get("reference_normalized_final") or r.get("reference_normalized_ours") or ""
        pred_norm = r.get("prediction_normalized") or ""
        if ref_norm:
            nc = len(ref_norm.replace(" ", ""))
            if nc > 0:
                cer_val = compute_cer([ref_norm], [pred_norm])
                total_char_err += cer_val / 100 * nc
                total_char_ref += nc

    ref1    = max(total_ref, 1)
    wer     = round(total_err  / ref1 * 100, 2)
    wer_std    = round(float(np.std(per_wers)),    2) if per_wers else 0.0
    wer_median = round(float(np.median(per_wers)), 2) if per_wers else 0.0
    sub_rate = round(total_subs / ref1 * 100, 2)
    ins_rate = round(total_ins  / ref1 * 100, 2)
    del_rate = round(total_dels / ref1 * 100, 2)
    cer      = round(total_char_err / max(total_char_ref, 1) * 100, 2)
    rtfx     = round(total_audio / total_time, 1) if total_time > 0 else None

    by_cat_out: dict[str, dict] = {}
    for cat, d in by_cat.items():
        cw     = round(d["err"] / max(d["ref"], 1) * 100, 2)
        std    = round(float(np.std(d["wers"])),    2) if d["wers"] else 0.0
        median = round(float(np.median(d["wers"])), 2) if d["wers"] else 0.0
        by_cat_out[cat] = {"wer": cw, "wer_std": std, "wer_median": median, "n": d["n"]}

    return {
        "wer": wer, "wer_std": wer_std, "wer_median": wer_median, "n": len(records),
        "cer": cer,
        "sub_rate": sub_rate, "ins_rate": ins_rate, "del_rate": del_rate,
        "rtfx": rtfx,
        "by_category": by_cat_out,
    }


def _build_leaderboard() -> list[dict]:
    # {model_id: {benchmark: {group: [records]}}}
    buckets: dict[str, dict] = defaultdict(lambda: {
        "public":   {"telephony": [], "domains": [], "short_inputs": [], "opensource": []},
        "private":  {"telephony": [], "domains": [], "short_inputs": [], "opensource": []},
        "combined": {"telephony": [], "domains": [], "short_inputs": [], "opensource": []},
    })

    for r in _all_records:
        mid = r["model_id"]
        bm  = r.get("_benchmark", "public")
        grp = _cat_group(r.get("category", ""))
        buckets[mid][bm][grp].append(r)
        buckets[mid]["combined"][grp].append(r)

    _GROUPS = ("telephony", "domains", "short_inputs", "opensource")

    rows = []
    for model_id, bm_data in buckets.items():
        row: dict = {"model_id": model_id}
        for bm in ("public", "private", "combined"):
            for grp in _GROUPS:
                stats = _compute_stats(bm_data[bm][grp])
                if stats:
                    row.setdefault(bm, {})[grp] = stats

        # Per-benchmark overall stats
        for bm in ("public", "private"):
            bm_recs = [r for grp in _GROUPS for r in bm_data[bm][grp]]
            if bm_recs:
                bm_overall = _compute_stats(bm_recs)
                if bm_overall:
                    row[f"{bm}_overall"] = bm_overall

        # Overall stats + rare WER — pooled across all categories, not split by group.
        all_recs = [r for grp in _GROUPS for r in bm_data["combined"][grp]]
        if all_recs:
            overall = _compute_stats(all_recs)
            if overall:
                row["overall"] = overall
            if _common_words:
                refs_all = [
                    r.get("reference_normalized_final") or r.get("reference_normalized_ours") or ""
                    for r in all_recs
                ]
                hyps_all = [r.get("prediction_normalized") or "" for r in all_recs]
                rare = compute_rare_wer(refs_all, hyps_all, _common_words)
                row["rare_wer"]      = rare["rare_wer"]
                row["rare_sub_rate"] = rare["rare_sub_rate"]

        rows.append(row)

    return rows


_leaderboard = _build_leaderboard()

# ─── Noise leaderboard ─────────────────────────────────────────────────────────


def _build_noise_leaderboard() -> list[dict]:
    if not _noise_tags:
        return []

    stats: dict[str, dict] = {}
    for model_id in _model_ids:
        stats[model_id] = {lv: {"errors": 0, "ref_words": 0, "n": 0} for lv in NOISE_LEVELS}

    # Build a combined sample map: public sample_map + private records
    combined: dict[tuple, dict[str, dict]] = dict(_pub_sample_map)
    for r in _all_records:
        if r.get("_benchmark") != "private":
            continue
        key = (r["split"], r["reference"], r.get("_ref_idx", 0))
        if key not in combined:
            combined[key] = {}
        combined[key][r["model_id"]] = r

    for (split, ref, _ref_idx), models in combined.items():
        noise_info = (_noise_tags.get(split) or {}).get(ref)
        if not noise_info:
            continue
        lv = noise_info["noise_level"]
        for model_id, r in models.items():
            if model_id not in stats:
                continue
            s  = r.get("word_substitutions", 0)
            i  = r.get("word_insertions", 0)
            d  = r.get("word_deletions", 0)
            nr = r.get("num_ref_words", 0)
            stats[model_id][lv]["errors"]    += s + i + d
            stats[model_id][lv]["ref_words"] += nr
            stats[model_id][lv]["n"]         += 1

    rows = []
    for model_id, by_level in stats.items():
        overall_errors = sum(v["errors"] for v in by_level.values())
        overall_ref    = sum(v["ref_words"] for v in by_level.values())
        overall_wer    = round(overall_errors / max(overall_ref, 1) * 100, 2)
        by_noise = {}
        for lv, d in by_level.items():
            if d["n"] == 0:
                continue
            wer = round(d["errors"] / max(d["ref_words"], 1) * 100, 2)
            by_noise[lv] = {"wer": wer, "n": d["n"]}
        rows.append({"model_id": model_id, "wer": overall_wer,
                     "n_samples": sum(v["n"] for v in by_level.values()), "by_noise": by_noise})

    rows.sort(key=lambda r: r["wer"])
    return rows


_noise_leaderboard = _build_noise_leaderboard()

# ─── Audio ─────────────────────────────────────────────────────────────────────

_ds_cache: dict[str, Any] = {}
_audio_idx: dict[str, dict[str, int]] = {}
_loading: set[str] = set()

PUBLIC_DATASET = "Revolab/ASR-Benchmark-Public"


def _ensure_index(split: str) -> bool:
    if split in _audio_idx:
        return True
    if split in _loading:
        return False
    _loading.add(split)
    try:
        from datasets import load_dataset
        logger.info("Building audio index for split=%s …", split)
        ds = load_dataset(PUBLIC_DATASET, split=split, streaming=False)
        _ds_cache[split] = ds
        _audio_idx[split] = {row["text"]: i for i, row in enumerate(ds.select_columns(["text"]))}
        logger.info("Audio index ready: %d entries (split=%s)", len(_audio_idx[split]), split)
        return True
    except Exception as exc:
        logger.warning("Cannot build audio index for split=%s: %s", split, exc)
        _audio_idx[split] = {}
        return False
    finally:
        _loading.discard(split)

# ─── Word confusion & deletions ────────────────────────────────────────────────

_word_confusion: dict = {}
_word_deletions: dict = {}


def _load_word_data() -> None:
    for d in (PUBLIC_DIR, LEGACY_DIR):
        cf = d / "word_confusion.json"
        df = d / "word_deletions.json"
        if cf.exists():
            with open(cf, encoding="utf-8") as f:
                _word_confusion.update(json.load(f))
            logger.info("Word confusion: %d models from %s", len(_word_confusion), d)
        if df.exists():
            with open(df, encoding="utf-8") as f:
                _word_deletions.update(json.load(f))
            logger.info("Word deletions loaded from %s", d)
        if cf.exists() or df.exists():
            break


_load_word_data()

# ─── Hard samples ──────────────────────────────────────────────────────────────

_hard_samples: list[dict] = []


def _build_hard_samples(min_wer: float = 50.0) -> list[dict]:
    rows = []
    for (split, ref, _ref_idx), models in _pub_sample_map.items():
        if not models:
            continue
        per_model = []
        for r in models.values():
            nref = r.get("num_ref_words", 0)
            if nref > 0:
                w = (r.get("word_substitutions", 0) + r.get("word_insertions", 0) + r.get("word_deletions", 0)) / nref * 100
                per_model.append({"model_id": r["model_id"], "wer": round(w, 1)})
        if not per_model or min(p["wer"] for p in per_model) < min_wer:
            continue
        first = next(iter(models.values()))
        noise_info = (_noise_tags.get(split) or {}).get(ref, {})
        rows.append({
            "ref":            ref,
            "ref_preview":    (ref[:80] + "…") if len(ref) > 80 else ref,
            "split":          split,
            "category":       first.get("category", ""),
            "avg_wer":        round(sum(p["wer"] for p in per_model) / len(per_model), 1),
            "min_wer":        round(min(p["wer"] for p in per_model), 1),
            "audio_length_s": round(first.get("audio_length_s", 0), 1),
            "noise_level":    noise_info.get("noise_level", ""),
            "models":         sorted(per_model, key=lambda x: x["wer"]),
        })
    rows.sort(key=lambda r: -r["avg_wer"])
    return rows


_hard_samples = _build_hard_samples()
logger.info("Hard samples (all models WER≥50%%): %d", len(_hard_samples))

# ─── API endpoints ─────────────────────────────────────────────────────────────


@app.get("/api/reload")
def api_reload():
    global _leaderboard, _noise_leaderboard, _hard_samples

    _all_records.clear()
    _pub_sample_map.clear()
    _pub_all_samples.clear()
    _model_ids.clear()
    _categories.clear()
    _splits.clear()
    _noise_tags.clear()
    _word_confusion.clear()
    _word_deletions.clear()
    _hard_samples.clear()

    _load_manifests()
    _load_noise_tags()
    _load_word_data()
    _leaderboard = _build_leaderboard()
    _noise_leaderboard = _build_noise_leaderboard()
    _hard_samples = _build_hard_samples()

    logger.info("Reloaded: %d models, %d public samples", len(_model_ids), len(_pub_all_samples))
    return {"status": "ok", "models": len(_model_ids), "samples": len(_pub_all_samples)}


@app.get("/api/meta")
def api_meta():
    has_private = bool(PRIVATE_DIR.exists() and any(PRIVATE_DIR.glob("*.jsonl")))
    return {
        "splits":        _splits,
        "categories":    _categories,
        "model_ids":     _model_ids,
        "total_samples": len(_pub_all_samples),
        "has_private":   has_private,
    }


@app.get("/api/leaderboard")
def api_leaderboard():
    return _leaderboard


@app.get("/api/leaderboard/noise")
def api_leaderboard_noise():
    return _noise_leaderboard


@app.get("/api/confusion")
def api_confusion():
    return _word_confusion


@app.get("/api/deletions")
def api_deletions():
    return _word_deletions


@app.get("/api/hard-samples")
def api_hard_samples(limit: int = 50):
    return {"samples": _hard_samples[:limit], "total": len(_hard_samples)}


@app.get("/api/analysis")
def api_analysis():
    for d in (PUBLIC_DIR, LEGACY_DIR):
        path = d / "model_analysis.json"
        if path.exists():
            with open(path, encoding="utf-8") as f:
                return json.load(f)
    return {}


@app.get("/api/samples")
def api_samples(split: str = "All", category: str = "All", noise: str = "All"):
    out = []
    for s in _pub_all_samples:
        if split not in ("All", "") and s["split"] != split:
            continue
        if category not in ("All", "") and s["category"] != category:
            continue
        ref = s["ref"]
        noise_info = (_noise_tags.get(s["split"]) or {}).get(ref, {})
        if noise not in ("All", "") and noise_info.get("noise_level", "") != noise:
            continue
        out.append({
            "ref":            ref,
            "ref_idx":        s.get("ref_idx", 0),
            "ref_preview":    (ref[:80] + "…") if len(ref) > 80 else ref,
            "split":          s["split"],
            "category":       s["category"],
            "best_wer":       s["best_wer"],
            "model_wers":     s.get("model_wers", {}),
            "audio_length_s": round(s["audio_length_s"], 1),
            "n_models":       s["n_models"],
            "noise_level":    noise_info.get("noise_level", ""),
            "snr_db":         noise_info.get("snr_db"),
        })
    return {"samples": out, "total": len(out)}


@app.get("/api/sample")
def api_sample(split: str, ref: str, ref_idx: int = 0):
    key = (split, ref, ref_idx)
    models_data = _pub_sample_map.get(key)
    if not models_data:
        raise HTTPException(404, "Sample not found")

    first = next(iter(models_data.values()))
    predictions = []
    for mid in _model_ids:
        if mid not in models_data:
            continue
        r = models_data[mid]
        ref_norm  = r.get("reference_normalized_final", "")
        pred_norm = r.get("prediction_normalized", "")
        subs = r.get("word_substitutions", 0)
        ins  = r.get("word_insertions", 0)
        dels = r.get("word_deletions", 0)
        num_ref = r.get("num_ref_words", 0)
        wer_val = round((subs + ins + dels) / max(num_ref, 1) * 100, 2) if num_ref > 0 else 0.0
        cer_val = compute_cer([ref_norm], [pred_norm]) if ref_norm else 0.0
        predictions.append({
            "model_id":                  mid,
            "prediction":                r.get("prediction", ""),
            "prediction_normalized":     pred_norm,
            "reference_normalized_final": ref_norm,
            "wer":                       wer_val,
            "cer":                       cer_val,
            "word_hits":                 r.get("word_hits", 0),
            "word_substitutions":        subs,
            "word_insertions":           ins,
            "word_deletions":            dels,
            "num_ref_words":             num_ref,
            "rtfx":                      round(r.get("rtfx", 0), 2),
            "transcription_time_ms":     round(r.get("transcription_time_s", 0) * 1000),
        })

    return {
        "reference":            first["reference"],
        "reference_normalized": first.get("reference_normalized_dataset_raw", first.get("reference_normalized_dataset", first.get("reference_normalized_final", ""))),
        "category":             first.get("category", ""),
        "audio_length_s":       round(first.get("audio_length_s", 0), 2),
        "split":                split,
        "predictions":          predictions,
    }


@app.get("/api/audio")
async def api_audio(request: Request, split: str, ref: str):
    ready = _ensure_index(split)
    if not ready and split not in _audio_idx:
        raise HTTPException(503, "Audio index building, please retry")

    idx = _audio_idx.get(split, {}).get(ref)
    if idx is None:
        raise HTTPException(404, "Audio not found for this sample")

    try:
        row = _ds_cache[split][idx]
        arr, sr = decode_audio(row["audio"], 16000)

        buf = io.BytesIO()
        try:
            import soundfile as sf
            sf.write(buf, arr, sr, format="WAV", subtype="PCM_16")
        except ImportError:
            import scipy.io.wavfile
            arr_i16 = (arr * 32767).clip(-32768, 32767).astype(np.int16)
            scipy.io.wavfile.write(buf, sr, arr_i16)

        body = buf.getvalue()
        total = len(body)

        # Mobile browsers (iOS Safari in particular) require proper HTTP Range
        # support for <audio> playback — without it, playback silently fails.
        range_header = request.headers.get("range")
        if range_header:
            try:
                units, _, range_spec = range_header.partition("=")
                start_s, _, end_s = range_spec.partition("-")
                start = int(start_s) if start_s else 0
                end = int(end_s) if end_s else total - 1
                end = min(end, total - 1)
            except ValueError:
                start, end = 0, total - 1
            chunk = body[start:end + 1]
            return Response(
                content=chunk,
                status_code=206,
                media_type="audio/wav",
                headers={
                    "Content-Range": f"bytes {start}-{end}/{total}",
                    "Accept-Ranges": "bytes",
                    "Content-Length": str(len(chunk)),
                    "Cache-Control": "public, max-age=3600",
                },
            )

        return Response(
            content=body,
            media_type="audio/wav",
            headers={
                "Accept-Ranges": "bytes",
                "Content-Length": str(total),
                "Cache-Control": "public, max-age=3600",
            },
        )
    except Exception as exc:
        logger.error("Audio error for split=%s idx=%d: %s", split, idx, exc)
        raise HTTPException(500, str(exc))


# ─── Serve frontend ────────────────────────────────────────────────────────────

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/")
def index():
    return FileResponse(str(STATIC_DIR / "index.html"))


if __name__ == "__main__":
    import uvicorn
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 9999
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")