"""Long-form inside the seat's envelope: every ~6-min window WHOLE, one request, per arm (arms one at a
time; arms run in parallel threads only when told to). Then score each window with the investigation's
gtscore.score against its slice of the timed ground truth, crosstalk split by boot.py's diarization rule.
usage:
  longwin.py run   ARM:PORT ...        -> out/raw/longwin/<arm>.jsonl  (text kept: public data)
  longwin.py score                     -> per arm: per-file WER (summed edits), dropouts, dropped words
"""
import glob
import json
import os
import sys
import time

AB = "/tank/spikes/parakeet-ab"
ROOT = "/tank/spikes/scriberr-slicer"
sys.path.insert(0, f"{AB}/code")
WINS = os.environ.get("AB_WINS", "long")   # long = placement 1, long2 = placement 2 (first cut at 180 s)
OUTD = f"{AB}/out/raw/longwin" + ("" if WINS == "long" else "-" + WINS)
os.makedirs(OUTD, exist_ok=True)
wins = [json.loads(l) for l in open(f"{AB}/data/{WINS}.jsonl")]

if sys.argv[1] == "run":
    import threading
    from bench import post

    def run(spec):
        arm, port = spec.split(":")
        with open(f"{OUTD}/{arm}.jsonl", "w") as fo:
            for w in wins:
                wav = open(w["wav"], "rb").read()
                t0 = time.time()
                try:
                    res = post(f"http://127.0.0.1:{port}/v1/audio/transcriptions", wav, timeout=1800)
                except Exception as e:  # noqa: BLE001
                    res = dict(status=-1, err=repr(e)[:300], e2e_ms=None, server_ms=None, text=None)
                fo.write(json.dumps(dict(arm=arm, id=w["id"], t0=t0, t1=time.time(), **res)) + "\n")
                fo.flush()
                print(arm, w["id"], res.get("status"), res.get("e2e_ms"), flush=True)
    ts = [threading.Thread(target=run, args=(s,)) for s in sys.argv[2:]]
    [t.start() for t in ts]
    [t.join() for t in ts]
else:
    from gtscore import score, tokens_from_text
    diar = {f: json.load(open(f"{ROOT}/public/{f}.diar.json"))["segments"] for f in ("wilde", "scotus")}
    gtt = {f: json.load(open(f"{ROOT}/gt/{f}.timed.json"))["times"] for f in ("wilde", "scotus")}
    for p in sorted(glob.glob(f"{OUTD}/*.jsonl")):
        res = {json.loads(l)["id"]: json.loads(l) for l in open(p)}
        arm = os.path.basename(p)[:-6]
        agg = {}
        for w in wins:
            r = res.get(w["id"])
            a = agg.setdefault(w["file"], dict(N=0, E=0, drops=0, dropw=0, clean=0, cleanw=0, ins=0, insw=0, fail=0, e2e=[]))
            if not r or r.get("status") != 200:
                a["fail"] += 1
                continue
            ref = w["ref_tokens"]
            s = score(ref, tokens_from_text(r["text"]))
            a["N"] += s["ref_words"]; a["E"] += s["S"] + s["D"] + s["I"]
            a["ins"] += s["insertion_runs"]; a["insw"] += s["insertion_words"]
            a["e2e"].append(r["e2e_ms"])
            # map window-relative token indices back to absolute GT times for the crosstalk rule
            times = [t for t in gtt[w["file"]] if w["t0"] <= t < w["t1"]]
            for g in s["drop_gaps"]:
                i0, i1 = g["ref_start"], g["ref_start"] + g["ref_words"]
                t0 = times[i0]
                t1 = times[i1] if i1 < len(times) else times[-1] + 0.5
                n = max(1, int((max(t1, t0 + 0.5) - t0) / 0.1))
                ov = sum(sum(1 for sg in diar[w["file"]] if sg["start"] <= t0 + k * 0.1 < sg["end"]) >= 2 for k in range(n)) / n
                a["drops"] += 1; a["dropw"] += g["ref_words"]
                if ov < 0.1:
                    a["clean"] += 1; a["cleanw"] += g["ref_words"]
        for f, a in agg.items():
            print(json.dumps(dict(arm=arm, file=f, windows=sum(1 for w in wins if w["file"] == f), failed=a["fail"],
                                  wer=round(100 * a["E"] / a["N"], 2) if a["N"] else None, ref_words=a["N"],
                                  dropouts=a["drops"], dropout_words=a["dropw"], clean_dropouts=a["clean"],
                                  clean_dropout_words=a["cleanw"], insertion_runs=a["ins"], insertion_words=a["insw"],
                                  e2e_ms_median=sorted(a["e2e"])[len(a["e2e"]) // 2] if a["e2e"] else None)))