#!/usr/bin/env python
# test_inference_parity.py — run_inference and infer_simple MUST produce identical output for the
# same policy (both drive the shared reset/step contract). Also checks the multi-emit contract:
# a frame whose step() returns a list emits several answers at that timestamp (the v5+v3 ensemble).
import importlib.util
import json
import os
import sys
import tempfile

from humomni.core.streaming_driver import run_sample, _emissions

# infer_simple lives in scripts/ (not a package) — load it by path so the test can call infer_one.
_p = os.path.join(os.path.dirname(__file__), "..", "scripts", "infer_simple.py")
_spec = importlib.util.spec_from_file_location("infer_simple", _p)
infer_simple = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(infer_simple)
infer_one = infer_simple.infer_one


class MockPolicy:
    """Deterministic, content-free: exercises None / str / list / [] returns from step()."""
    def reset(self, question):
        self.n = 0

    def step(self, t, frame_path):
        self.n += 1
        if self.n == 2:
            return "single"            # one answer (str)
        if self.n == 3:
            return ["a", "b"]          # TWO answers at one timestamp (the ensemble case)
        if self.n == 4:
            return []                  # explicit silence
        return None                    # silence


def _make_sample(d, n_frames=5):
    for i in range(1, n_frames + 1):
        open(os.path.join(d, f"{i * 0.5:.1f}.jpg"), "wb").close()   # empty frames on a 0.5 grid
    json.dump({"question_id": "vid.mp4", "question": "what happens?"},
              open(os.path.join(d, "question.json"), "w", encoding="utf-8"))


def test_parity_and_multiemit():
    with tempfile.TemporaryDirectory() as d:
        _make_sample(d)
        a = run_sample(d, MockPolicy())                 # run_inference path (streaming_driver.drive)
        b = infer_one(d, MockPolicy())                  # infer_simple path
        assert a == b, f"run_inference != infer_simple:\n  {a}\n  {b}"
        resp = a["model_response_list"]
        # frame 2 -> "single" @1.0; frame 3 -> "a","b" @1.5; frames 4,5 -> silent
        assert [r["content"] for r in resp] == ["single", "a", "b"], resp
        assert [r["time"] for r in resp] == [1.0, 1.5, 1.5], resp     # two emits share t=1.5
    # _emissions normalization
    assert _emissions(None) == [] and _emissions("x") == ["x"] and _emissions(["a", "b"]) == ["a", "b"]
    print("PASS: run_inference == infer_simple; multi-emit (ensemble) lands N answers at one timestamp.")


if __name__ == "__main__":
    sys.exit(test_parity_and_multiemit())