#!/usr/bin/env python3
"""
replay.py — replay a dumped embedding HDF5 through the real KPN downstream nodes.

TRACES: VR-002 | PR-002

Reads an embedding dump (scripts/optimizer/SCHEMA.md), feeds each frame as an
EmbeddedSceneFrame into a Python-assembled KPN network wiring the *real* C++
face_tracker → identity_matcher → frame_annotation, and returns the same presence-window
JSON that scene_analyze's result_sink produces (minimal schema). No decode, no GPU
embedding — only the cheap downstream tail runs, so a sweep can vary Config knobs
freely. See [[kpn-python-replay-optimizer]].

CLI:
  python scripts/optimizer/replay.py --dump film.h5 --gallery gallery.json \
      --out replayed.json [--prob-threshold 0.99] [--anneal-sec 10] ...
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import h5py
import numpy as np

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts"))
from sae_stamp import verify_gallery_stamp  # noqa: E402


def dump_embedder_stamp(dump_path: str) -> dict:
    """The GR-004 embedder stamp recorded in an embedding dump.

    A replay has no live embedder — the dump IS the embedder as far as the gallery
    is concerned, so the dump's stamp is what the gallery must be checked against.
    Dumps written before GR-004 have no attributes and yield an empty stamp, which
    the check reports as unverifiable rather than silently accepting."""
    with h5py.File(dump_path, "r") as f:
        name = f.attrs.get("embedder_model", "")
        sha = f.attrs.get("embedder_sha256", "")
    dec = lambda v: v.decode() if isinstance(v, bytes) else ("" if v is None else str(v))
    return {"model_name": dec(name), "model_sha256": dec(sha), "embed_dim": 512}


def load_frames(dump_path: str, min_conf: float = 0.0):
    """Yield EmbeddedSceneFrame dicts from the HDF5 dump, then a trailing EOF.

    `min_conf` drops detections below that detector confidence before they reach the
    matcher — an UPWARD-only detector_conf sweep on already-dumped faces (the dump was
    made at detector_conf=0.5, so 0.5 is the floor). Lets us test whether near-threshold
    detections are real faces (raising min_conf hurts recall) or phantoms (it helps
    precision at no recall cost)."""
    with h5py.File(dump_path, "r") as f:
        ts   = f["frames/timestamp_sec"][:]
        fidx = f["frames/frame_idx"][:]
        cut  = f["frames/is_cut"][:]
        off  = f["frames/face_offset"][:]
        cnt  = f["frames/face_count"][:]
        emb  = f["faces/embedding"][:]
        bbox = f["faces/bbox"][:]
        lmk  = f["faces/landmarks"][:]
        conf = f["faces/confidence"][:]
        # TRACES: AR-028 | SR-002
        # The quality vector, present from schema v2. A v1 dump predates AR-028
        # and simply has no such dataset — read as absent, never as a default,
        # so a face from an old dump stays at the C++ -1 "unscored" sentinel
        # rather than acquiring a fabricated sharpness of 0 (which is a real
        # value on this axis, meaning a featureless crop).
        qual = {k: f[f"faces/{k}"][:] for k in ("sharpness", "alignment_residual")
                if f"faces/{k}" in f}
        movie = f.attrs.get("movie", "")
        fps   = float(f.attrs.get("sample_fps", 1.0))

    frames = []
    for i in range(len(ts)):
        s, n = int(off[i]), int(cnt[i])
        keep = slice(s, s + n)
        c = np.ascontiguousarray(conf[keep], dtype=np.float32)
        if min_conf > 0.0 and n:
            m = c >= min_conf
            sel = np.where(m)[0]
            frames.append({
                "timestamp_sec": float(ts[i]), "frame_idx": int(fidx[i]),
                "is_cut": bool(cut[i]), "eof": False,
                "bbox":       np.ascontiguousarray(bbox[keep][sel], dtype=np.float32),
                "landmarks":  np.ascontiguousarray(lmk[keep][sel],  dtype=np.float32),
                "confidence": np.ascontiguousarray(c[sel],          dtype=np.float32),
                "embeddings": np.ascontiguousarray(emb[keep][sel],  dtype=np.float32),
                **{k: np.ascontiguousarray(v[keep][sel], dtype=np.float32)
                   for k, v in qual.items()},
            })
        else:
            frames.append({
                "timestamp_sec": float(ts[i]), "frame_idx": int(fidx[i]),
                "is_cut": bool(cut[i]), "eof": False,
                "bbox":       np.ascontiguousarray(bbox[keep], dtype=np.float32),
                "landmarks":  np.ascontiguousarray(lmk[keep],  dtype=np.float32),
                "confidence": c,
                "embeddings": np.ascontiguousarray(emb[keep],  dtype=np.float32),
                **{k: np.ascontiguousarray(v[keep], dtype=np.float32)
                   for k, v in qual.items()},
            })
    last_ts = float(ts[-1]) if len(ts) else 0.0
    frames.append({"timestamp_sec": last_ts, "eof": True})
    return frames, str(movie), fps


def replay(dump_path: str, gallery: str, cfg: dict, build_dir: str, stop: bool = True,
          raw_out: str | None = None) -> dict:
    """Run the dump through the real KPN chain; return minimal-schema presence JSON.

    cfg may include "detector_conf" to prune dumped detections below that confidence
    (upward-only from the 0.5 dump floor) before matching.

    raw_out: if set, also write the raw per-frame annotations (timestamp, actor_idx,
    name, bbox, similarity — one entry per input frame, before merging into windows)
    as JSON lines to this path. Needed to draw bounding boxes on extracted frames;
    the merged window schema returned by this function has no per-frame bbox."""
    sys.path.insert(0, build_dir)
    import sae_kpn

    # TRACES: GR-004 | SR-001
    # checked here, before any network is built, so a
    # cross-model replay dies with one readable error instead of producing a
    # plausible-looking score. add_identity_matcher re-checks it C++-side below;
    # that is the backstop for any other caller of the binding.
    stamp = dump_embedder_stamp(dump_path)
    verify_gallery_stamp(gallery, stamp=stamp,
                         embedder_desc=f"embedding dump {Path(dump_path).name}",
                         require_stamp=bool(cfg.get("require_gallery_stamp", False)))

    frames, movie, fps = load_frames(dump_path, min_conf=float(cfg.get("detector_conf", 0.0)))

    net = sae_kpn.Network()
    sae_kpn._register_types(net)

    idx = [0]
    eof = {"timestamp_sec": frames[-1]["timestamp_sec"], "eof": True}

    def source():
        # A no-input source node's run_loop calls this in a tight loop. Once frames
        # are exhausted we must NOT hot-spin returning EOF — that pegs a core and
        # floods the downstream channel with EOFs (livelock that wedged DE). Sleep
        # briefly after the single real EOF so net.stop() can tear the thread down.
        i = idx[0]
        idx[0] += 1
        if i < len(frames):
            return frames[i]
        time.sleep(0.05)
        return eof

    # Channel capacity must exceed the frame count so the fast source can't overflow
    # a downstream FIFO before the serial reader drains it — PyNode DROPS on overflow,
    # which would silently truncate the replay. Size to the whole film + slack.
    # Every channel gets capacity ≥ the whole film so NOTHING can ever overflow-drop:
    # the source can push all frames before any downstream node has drained, and a
    # dropped frame silently corrupts the score. Memory is cheap (a few k pointers);
    # correctness is not. Generous slack on top.
    cap = len(frames) * 2 + 64
    sae_kpn.add_node_python(net, "replay", source, [], ["EmbeddedSceneFrame"], cap)
    sae_kpn.add_face_tracker(net, "tracker", cfg, cap)
    sae_kpn.add_identity_matcher(net, "matcher", gallery, cfg, cap,
                                 stamp["model_name"], stamp["model_sha256"])
    sae_kpn.add_frame_annotation(net, "scene", cap)
    net.connect("replay", 0, "tracker", 0)
    net.connect("tracker", 0, "matcher", 0)
    net.connect("matcher", 0, "scene", 0)
    net.build()
    net.start()

    # Read exactly one annotation per input frame. The source emits EOF as an ordinary
    # value AFTER the last frame, but the concurrent pipeline lets that EOF OVERTAKE
    # the last few real frames still flowing tracker→matcher→scene. Breaking on the
    # first eof therefore dropped a random tail (~0.5–1%, race-dependent). Instead we
    # keep reading past eof until we've collected all n_frames annotations (or hit a
    # run of consecutive eofs meaning the pipeline is genuinely drained).
    n_expected = len(frames) - 1   # excludes the trailing eof frame
    annotations = []
    eof_streak = 0
    max_reads = n_expected * 2 + 32
    for _ in range(max_reads):
        sa = net.read("scene", 0)
        if sa.get("eof"):
            eof_streak += 1
            # stragglers can still arrive after an eof; only stop once we've either
            # got everything or seen several eofs in a row (truly drained).
            if len(annotations) >= n_expected or eof_streak >= 8:
                break
            continue
        eof_streak = 0
        annotations.append(sa)
        if len(annotations) >= n_expected:
            break

    if raw_out:
        with open(raw_out, "w") as f:
            for sa in annotations:
                f.write(json.dumps(sa) + "\n")

    result = build_minimal(annotations, movie, fps, cfg)
    if stop:
        net.stop()
    return result


def build_minimal(annotations, movie, fps, cfg) -> dict:
    """Per-actor [start,end] windows, built by annealing per-frame detections.

    TRACES: VR-011 | PR-002

    This NO LONGER mirrors ResultSinkFunc, and the docstring used to claim it
    did. The sink builds a window from a TrackRegistry claim -- the extent
    [first_seen, last_seen] of a track an actor owned (AR-012) -- so a window
    starts when the actor appeared rather than when recognition first
    succeeded, and interior gaps are absorbed by the track surviving them.
    This function still bridges gaps between isolated accepted frames, which is
    what anneal_sec did before AR-012/AR-013 withdrew it.

    So a replayed window and a pipeline window are answers to different
    questions, and a sweep tuned against this one is not tuning the shipped
    behaviour. That is VR-011's job -- "rewrite the replay harness for the
    post-AR-012 output contract" -- and it is a rewrite, not an edit, because
    the registry's claims do not cross the Python seam at all today.

    `anneal_sec` is therefore replay-local now: it configures THIS function and
    is no longer forwarded to the C++ Config, which has no such field.
    """
    anneal = float(cfg.get("anneal_sec", 10.0))
    info = {}           # actor_idx -> identity fields
    times = {}          # actor_idx -> [timestamps]
    for sa in annotations:
        for a in sa["visible_actors"]:
            if a["actor_idx"] < 0:
                continue
            info[a["actor_idx"]] = a
            times.setdefault(a["actor_idx"], []).append(sa["timestamp_sec"])

    actors = []
    for idx, ts in times.items():
        ts.sort()
        scenes = []
        ws = we = ts[0]
        for t in ts[1:]:
            if t - we > anneal:
                scenes.append([ws, we])
                ws = t
            we = t
        scenes.append([ws, we])
        a = info[idx]
        actors.append({
            "name": a["name"], "imdb_id": a["imdb_id"], "tmdb_id": a["tmdb_id"],
            "jellyfin_id": a["jellyfin_id"], "scenes": scenes,
        })

    return {"schema_version": 1, "movie": movie, "sample_fps": fps,
            "anneal_sec": anneal, "actors": actors}


# TRACES: AR-024 | SR-002
# Keys the C++ Config actually still has. Seven names were removed here, all of
# them accepted silently for months after the fields behind them were deleted:
#
#   match_threshold, match_ratio, match_ratio_ceil  — the raw-cosine accept
#     fallback, retired with AR-024's enforcement.
#   track_max_embed_dist, cut_revive_sim            — raw cosines, retired
#     earlier by AR-024 when association moved into probability space.
#   track_max_frames_missing, cut_inactive_max_frames — frame counts whose
#     meaning changed with sample_fps, retired by AR-008/AR-013 in favour of
#     track_extinction_sec.
#
# A sweep that varied one of these was measuring nothing, and reported a
# perfectly ordinary-looking F1 for its trouble. kpn_bindings.cpp reads config
# keys with a contains() check, so an unknown key is not an error — which makes
# a stale entry here silently inert rather than loudly wrong.
CFG_KEYS = ["detector_conf", "prob_threshold", "match_prior",
            "track_alpha", "track_min_iou", "track_assoc_min_prob",
            "track_extinction_sec",
            # AR-025 ownership and evidence accumulation. Newly reachable:
            # these were in-class defaults no sweep could vary, which is why
            # VR-007 never covered them despite rho_max deferring to it.
            "ownership_logodds", "evidence_rho_max", "evidence_admit_below",
            "evidence_max_views"]

# Swept like a Config key but consumed entirely in Python, by build_minimal.
# Kept separate so nobody has to guess which of these the pipeline actually
# reads: everything in CFG_KEYS crosses the seam, and nothing here does.
REPLAY_LOCAL_KEYS = ["anneal_sec"]


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--dump", required=True, help="embedding HDF5 dump")
    p.add_argument("--gallery", required=True)
    p.add_argument("--out", required=True, help="output presence JSON")
    p.add_argument("--raw-out", help="also write raw per-frame annotations (JSONL, with bboxes) here")
    p.add_argument("--build-dir", default=str(REPO / "build"))
    for k in CFG_KEYS + REPLAY_LOCAL_KEYS:
        p.add_argument(f"--{k.replace('_','-')}", type=float, default=None)
    # per-film gallery expansion: promotes pose-varied views of confidently-identified
    # actors into an in-memory annex, recovering ~+4 recall at no precision cost.
    p.add_argument("--expand-gallery", action="store_true")
    # TRACES: GR-004 | SR-001
    # promote an unprovable gallery/dump binding from a
    # loud warning to a hard error. Measurement sweeps should set this (or
    # SAE_REQUIRE_GALLERY_STAMP=1) so no number comes from an unbound pair.
    p.add_argument("--require-gallery-stamp", action="store_true")
    args = p.parse_args()

    # Both lists go into one dict: config_from_dict reads C++ keys with a
    # contains() check and ignores the rest, and build_minimal reads its own.
    cfg = {k: getattr(args, k)
           for k in CFG_KEYS + REPLAY_LOCAL_KEYS if getattr(args, k) is not None}
    if args.expand_gallery:
        cfg["expand_gallery"] = True
    if args.require_gallery_stamp:
        cfg["require_gallery_stamp"] = True
    # stop=True: PyNode::stop() sets stop_flag_ before joining, so the source
    # thread's run_loop actually exits. stop=False skips that, leaving stop_flag_
    # false forever — the PyNode destructor's jthread.join() then blocks forever
    # (verified via gdb: stuck in the source node's run_loop, not the GEMM path).
    result = replay(args.dump, args.gallery, cfg, args.build_dir, stop=True,
                    raw_out=args.raw_out)
    Path(args.out).write_text(json.dumps(result, indent=2))
    print(f"[replay] {len(result['actors'])} actors → {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()