#!/usr/bin/env python3
"""
Smoke test for the sae_kpn module: assemble the real downstream pipeline
(tracker → matcher → annotation → sink) in a Python-driven KPN network, fed by a
no-input Python source node, and verify the sink writes a truth file.

TRACES: VR-011 | PR-002

Proves the KPN-native replay path works without any numpy port of node logic.

Rewritten for `add_pipeline`. It previously called three node factories and read
SceneAnnotations back through the seam, asserting on what came out per frame.
Neither half of that survives VR-011: the factories are gone because the chain
has a construction order Python could not express, and presence is now the C++
sink's answer, derived from TrackRegistry claims. Nothing is read per frame, so
the assertions are on the file the sink writes.

Run:  python scripts/optimizer/test_sae_kpn.py [gallery.json] [build_dir]
"""
import json
import sys
import tempfile
import time
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parent.parent.parent
GAL = sys.argv[1] if len(sys.argv) > 1 else str(REPO / "gallery_arcface_w600k_r50.json")
BUILD = sys.argv[2] if len(sys.argv) > 2 else str(REPO / "build")
sys.path.insert(0, BUILD)
import sae_kpn  # noqa: E402


def make_frame(t, n):
    e = np.random.randn(n, 512).astype(np.float32)
    e /= np.linalg.norm(e, axis=1, keepdims=True)
    return {"timestamp_sec": t, "eof": False,
            "bbox": np.tile(np.array([10, 10, 50, 50], np.float32), (n, 1)),
            "landmarks": np.tile(np.arange(10, dtype=np.float32), (n, 1)),
            "confidence": np.full((n,), 0.9, np.float32), "embeddings": e}


def main():
    net = sae_kpn.Network()
    sae_kpn._register_types(net)

    frames = [make_frame(float(t), 1) for t in range(3)]
    frames.append({"timestamp_sec": 3.0, "eof": True})
    idx = [0]
    eof_frame = {"timestamp_sec": 3.0, "eof": True}

    def source():
        # Emit each frame once, then keep returning EOF so the node thread stays
        # responsive to stop(). The sleep matters: a no-input source is called in
        # a tight loop, and hot-spinning EOFs pegs a core and floods the channel.
        i = idx[0]
        idx[0] += 1
        if i < len(frames):
            return frames[i]
        time.sleep(0.05)
        return eof_frame

    with tempfile.TemporaryDirectory() as tmp:
        out_path = str(Path(tmp) / "truth.json")
        cfg = {
            "prob_threshold": 0.99,
            "track_extinction_sec": 5.0,
            "output_path": out_path,
            "movie_path": "sae_kpn smoke test",
            "sample_fps": 1.0,
            # Standard verbosity emits the per-frame array this test asserts on.
            # At 0 the file carries only the actor epochs, and three random
            # embeddings against a real gallery need not produce any.
            "verbosity": 1,
        }

        sae_kpn.add_node_python(net, "replay", source, [], ["EmbeddedSceneFrame"], 16)
        # No embedder stamp: these embeddings are random, not the output of any
        # model, so there is nothing truthful to claim. That warns rather than
        # failing, and would be fatal under SAE_REQUIRE_GALLERY_STAMP — which is
        # correct, since an unverifiable binding is exactly what it guards.
        sae_kpn.add_pipeline(net, GAL, cfg, 16)

        net.connect("replay", 0, "tracker", 0)
        net.connect("tracker", 0, "matcher", 0)
        net.connect("matcher", 0, "annotation", 0)
        net.connect("annotation", 0, "sink", 0)
        net.build()
        net.start()

        # The sink writes on the EOF annotation. Wait for that rather than
        # reading anything back: presence lives entirely on the C++ side.
        deadline = time.time() + 30.0
        while not sae_kpn.pipeline_done(net):
            if time.time() > deadline:
                sae_kpn.release_pipeline(net)
                raise TimeoutError("sink never saw EOF within 30s")
            time.sleep(0.02)

        net.stop()
        sae_kpn.release_pipeline(net)

        with open(out_path) as f:
            truth = json.load(f)

    per_frame = truth.get("frames", [])
    assert "actors" in truth, "truth file has no actors array"
    assert len(per_frame) == 3, f"expected 3 frames, got {len(per_frame)}"
    # EOF is a control token, not an observation: the sink flushes on it and does
    # not record it, so three inputs give three frames and never four.
    assert [f["t"] for f in per_frame] == [0.0, 1.0, 2.0], "timestamps wrong"
    assert all("identified" in f for f in per_frame), "missing identified"
    print(f"OK: {len(per_frame)} frames through the real KPN chain, sink wrote its truth file")


if __name__ == "__main__":
    main()