import json
from types import SimpleNamespace

import mlx.core as mx

from h3_mlx import benchmark, runtime


def test_summary_excludes_warmup_and_retains_every_resident_sample(
    tmp_path, monkeypatch
):
    calls = []
    cleared = []

    class FakePipeline:
        def __init__(self, *args, **kwargs):
            assert kwargs["kernel"] == "metal-qk-exact"
            assert kwargs["tune_matmul"] is True

        def generate(self, prompt, **kwargs):
            calls.append(kwargs)
            return None, {
                "png_sha256": "a" * 64,
                "conditioning_cache_hit": len(calls) > 1,
                "dit_was_resident": len(calls) > 1,
            }

        def clear(self):
            cleared.append(True)

    clock = iter([0.0, 50.0, 50.0, 55.0, 55.0, 62.0])
    monkeypatch.setattr(
        benchmark, "time", SimpleNamespace(perf_counter=lambda: next(clock))
    )
    monkeypatch.setattr(runtime, "Pipeline", FakePipeline)
    monkeypatch.setattr(runtime, "environment", lambda: {"machine": "test"})
    monkeypatch.setattr(mx, "synchronize", lambda: None)
    output = tmp_path / "run"
    benchmark.main(
        [
            "--model",
            "model",
            "--text-encoder",
            "encoder",
            "--prompt",
            "test",
            "--warmup",
            "1",
            "--repeats",
            "2",
            "--output",
            str(output),
        ]
    )
    summary = json.loads((output / "summary.json").read_text())
    assert summary["median_seconds"] == 6.0
    assert [s["wall_seconds"] for s in summary["samples"]] == [50.0, 5.0, 7.0]
    assert [s["phase"] for s in summary["samples"]] == [
        "warmup",
        "measured",
        "measured",
    ]
    assert all(s["dit_was_resident"] for s in summary["samples"][1:])
    assert all(c["save_trace"] for c in calls)
    assert cleared == [True]