"""Tests for the milestone-gqa-headline bench.

The bench wires the new ``_gqa_attention_prefill_long`` kernel into a
single-KV-group milestone panel (LLaMA-3.1-70B target, C=8 P=8). The
4 legacy C=1 / C=4 panels (single_user_*, multi_user_*) were removed
once pytest regression covered those configurations comprehensively
and the comparative decode work moved into the
``milestone-gqa-decode-4cases`` bench.

Single SIP scope (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
headline deferred per ADR-0060 §B-item-1).
"""
from __future__ import annotations

import json

import kernbench.benches.milestone_gqa_headline as bench_mod  # noqa: F401 (Phase 2)
from kernbench.benches.registry import resolve
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology

BENCH_NAME = "milestone-gqa-headline"

PANELS = (
    "single_kv_group_prefill_gqa_c8_p8",
)


def _run_validation():
    topo = resolve_topology("topology.yaml")
    return run_bench(
        topology=topo,
        bench_fn=resolve(BENCH_NAME).run,
        device=resolve_device(None),
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )


def _sweep_json(monkeypatch) -> dict:
    monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
    out = bench_mod._OUTPUT_DIR / "sweep.json"
    if not out.exists():
        result = _run_validation()
        assert result.completion.ok, result.completion
    assert out.exists(), f"missing {out}"
    return json.loads(out.read_text())


# ── Registration ───────────────────────────────────────────────────────


def test_bench_registered():
    spec = resolve(BENCH_NAME)
    assert spec.name == BENCH_NAME
    assert callable(spec.run)
    assert spec.description.strip(), "description must be non-empty"


# ── Validation run completes ──────────────────────────────────────────


def test_validation_run_completes_ok(monkeypatch):
    monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
    result = _run_validation()
    assert result.completion.ok, (
        f"headline validation run failed: {result.completion}"
    )


# ── sweep.json shape ──────────────────────────────────────────────────


def test_sweep_json_has_expected_panels(monkeypatch):
    data = _sweep_json(monkeypatch)
    assert set(data["panels"]) == set(PANELS), (
        f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}"
    )
    assert len(data["rows"]) == len(PANELS)
    assert {r["panel"] for r in data["rows"]} == set(PANELS)


# Per-panel architectural assertions for the surviving single-KV-group
# panel live in tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py.
# Decode architectural assertions live in
# tests/attention/test_milestone_gqa_decode_4cases.py.