#!/usr/bin/env python3
"""
Smoke test for the sae_kpn module: assemble the real downstream pipeline nodes
(face_tracker → identity_matcher → scene_tracker) in a Python-driven KPN network,
fed by a no-input Python source node, and verify SceneAnnotations flow out.

Proves the KPN-native replay path works without any numpy port of node logic.
Run:  python scripts/optimizer/test_sae_kpn.py [gallery.json] [build_dir]
"""
import sys
import queue
import numpy as np
from pathlib import Path

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)
    cfg = {"prob_threshold": 0.99, "anneal_sec": 10.0, "extinction_sec": 5.0}

    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 (never block) so the node
        # thread stays responsive to stop() after the sink has seen EOF.
        i = idx[0]
        idx[0] += 1
        return frames[i] if i < len(frames) else eof_frame

    sae_kpn.add_node_python(net, "replay", source, [], ["EmbeddedSceneFrame"], 8)
    sae_kpn.add_face_tracker(net, "tracker", cfg, 16)
    sae_kpn.add_identity_matcher(net, "matcher", GAL, cfg, 16)
    sae_kpn.add_scene_tracker(net, "scene", cfg, 16)
    net.connect("replay", 0, "tracker", 0)
    net.connect("tracker", 0, "matcher", 0)
    net.connect("matcher", 0, "scene", 0)
    net.build()
    net.start()

    got = []
    for _ in range(4):
        sa = net.read("scene", 0)
        got.append(sa)
        if sa.get("eof"):
            break
    net.stop()

    non_eof = [g for g in got if not g.get("eof")]
    assert len(non_eof) == 3, f"expected 3 annotations, got {len(non_eof)}"
    assert got[-1].get("eof"), "expected trailing EOF"
    assert [g["timestamp_sec"] for g in non_eof] == [0.0, 1.0, 2.0], "timestamps wrong"
    assert all("visible_actors" in g for g in non_eof), "missing visible_actors"
    print(f"OK: {len(non_eof)} annotations through the real KPN chain, EOF received")


if __name__ == "__main__":
    main()