"""Tests for the compute-bound prefill composite-command study.

Unlike memory-bound decode (where command form is latency-neutral), in
compute-bound prefill the composite command keeps the MAC array fed by
pipelining DMA↔compute per HW tile, so it wins on wall-clock latency. The
rigorous claim here is therefore the *opposite* of the decode study:
``composite_latency < primitive_latency`` at a compute-bound context.
"""
from __future__ import annotations

from pathlib import Path

from kernbench.common.pe_commands import CompositeCmd
from kernbench.triton_emu.tl_context import TLContext, run_kernel

from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
    gqa_prefill_composite_ext_kernel,
    gqa_prefill_composite_kernel,
    gqa_prefill_primitive_kernel,
)

_TOPOLOGY = Path(__file__).resolve().parents[2] / "topology.yaml"
_H_Q, _H_KV, _D_HEAD = 8, 1, 128


def _n_composites(kernel, T_q: int, S_kv: int) -> int:
    tl = TLContext(pe_id=0, num_programs=1, scratch_base=1 << 61,
                   scratch_size=1 << 20)
    run_kernel(kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
               T_q, S_kv, _H_Q, _H_KV, _D_HEAD, 1, 1)
    return sum(1 for c in tl.commands if isinstance(c, CompositeCmd))


def test_command_form_composite_counts():
    """One Q-block (T_q=16 → M=128): the primitive issues no composites;
    each composite form issues exactly two (Q·Kᵀ and P·V / recipe)."""
    assert _n_composites(gqa_prefill_primitive_kernel, 16, 256) == 0
    assert _n_composites(gqa_prefill_composite_kernel, 16, 256) == 2
    assert _n_composites(gqa_prefill_composite_ext_kernel, 16, 256) == 2


def _latency_ns(variant: str, ctx_len: int) -> float:
    from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
        _run_panel,
    )
    return _run_panel(variant, ctx_len, str(_TOPOLOGY))["latency_ns"]


def test_three_variants_complete_in_data_mode():
    """All three prefill forms run end-to-end through the engine."""
    for variant in ("primitive", "composite", "composite_extended"):
        assert _latency_ns(variant, 256) > 0


def test_composite_faster_in_compute_bound_prefill():
    """At a compute-bound context (512), the composite command's DMA↔compute
    pipelining beats the primitive's serial load→dot — the opposite of the
    memory-bound decode study, where the two tie."""
    prim = _latency_ns("primitive", 512)
    comp = _latency_ns("composite", 512)
    assert comp < prim, f"composite {comp} not < primitive {prim}"