#!/usr/bin/env python3
"""Run person detection over a whole clip and cache it to detections.jsonl.

Detection is the only expensive stage, so it is isolated here and written once.
Everything downstream (tracking, zones, cycles, metrics) replays from the cache
in seconds -- see build_detections_cache.py for why that mattered.

Supply several `--ports` to fan out across parallel detector services; the
service is a single-threaded HTTP handler, so one process saturates at one
request in flight. Resumes automatically: an existing detections.jsonl with the
expected frame count is left alone, and a partial one is topped up.
"""

from __future__ import annotations

import argparse
import json
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path


def detect(port, frame_path, frame_id, run_id, timeout=120.0):
    import urllib.request

    payload = {
        "schema_version": "roven.edge.detector_service_request.v0",
        "run_id": run_id,
        "model_profile": "ms1_live_detector",
        "frames": [
            {
                "image": frame_id,
                "path": str(frame_path),
                "camera_id": "tempering_cam",
                "time": time.strftime("%Y-%m-%dT%H:%M:%S"),
                "frame_ref": f"file://{frame_path}",
            }
        ],
    }
    req = urllib.request.Request(
        f"http://127.0.0.1:{port}/roven/detect",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    dets = []
    for record in result.get("records", []):
        for item in record.get("detections", []):
            label = str(item.get("label") or "").lower()
            if label and label != "person":
                continue
            bbox = [int(round(float(v))) for v in (item.get("bbox_xyxy") or [])[:4]]
            if len(bbox) == 4:
                dets.append({"bbox_xyxy": bbox, "confidence": round(float(item.get("confidence") or 0.0), 4)})
    return dets


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--frames-dir", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--ports", default="18085")
    ap.add_argument("--video-t0", type=int, default=0)
    ap.add_argument("--run-id", default="glass-batch")
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    ports = [int(p) for p in args.ports.split(",") if p.strip()]
    frames = sorted(Path(args.frames_dir).glob("f_*.jpg"))
    if args.limit:
        frames = frames[: args.limit]
    n = len(frames)
    out_path = Path(args.out)

    done = {}
    for src in (out_path, out_path.with_suffix(out_path.suffix + ".part")):
        if not src.exists():
            continue
        with src.open(encoding="utf-8") as fh:
            for line in fh:
                try:
                    r = json.loads(line)
                    done[r["t_local"]] = r
                except json.JSONDecodeError:
                    pass  # a partial last line from a killed run
    if out_path.exists() and len(done) >= n:
        print(f"{args.out}: already complete ({len(done)} frames)")
        return 0
    if done:
        print(f"resuming: {len(done)}/{n} already done")

    todo = [(i, p) for i, p in enumerate(frames) if i not in done]
    t0 = time.time()
    results = dict(done)
    errors = 0

    def work(job):
        i, p = job
        port = ports[i % len(ports)]
        try:
            d = detect(port, p.resolve(), f"{args.run_id}_{i:06d}", args.run_id)
        except Exception as exc:  # noqa: BLE001
            return i, None, str(exc)
        return i, d, None

    # Append every result as it lands. A 4 h clip is ~2 h of detection on CPU;
    # buffering it all in memory and writing once at the end meant any restart
    # threw the whole run away. The file is unordered while in progress and
    # gets sorted on completion, and resume just reads whatever is there.
    part = out_path.with_suffix(out_path.suffix + ".part")
    with ThreadPoolExecutor(max_workers=len(ports)) as ex, part.open("a", encoding="utf-8") as pf:
        for k, (i, dets, err) in enumerate(ex.map(work, todo), 1):
            if err is not None:
                errors += 1
                if errors <= 5:
                    print(f"  [t={i}] {err}", flush=True)
                dets = []
            rec = {"t_local": i, "t_abs": args.video_t0 + i, "detections": dets}
            results[i] = rec
            pf.write(json.dumps(rec, ensure_ascii=False) + "\n")
            if k % 200 == 0:
                pf.flush()
            if k % 500 == 0 or k == len(todo):
                el = time.time() - t0
                rate = k / max(el, 1e-9)
                print(
                    f"  {k}/{len(todo)}  {rate:.1f} f/s  eta {(len(todo)-k)/max(rate,1e-9)/60:.1f} min",
                    flush=True,
                )

    with out_path.open("w", encoding="utf-8") as fh:
        for i in range(n):
            r = results.get(i) or {"t_local": i, "t_abs": args.video_t0 + i, "detections": []}
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    part.unlink(missing_ok=True)
    tot = sum(len(r["detections"]) for r in results.values())
    print(f"{args.out}: {n} frames, {tot} detections, {errors} errors, {(time.time()-t0)/60:.1f} min")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())