"""Command-count checks for the Case-6 primitive-TILED (streamed) decode kernel.

The tiled kernel hand-blocks each Q·Kᵀ / P·V matmul into 16×16×16
GEMMs, with per-block DMAs of the HBM operand slices and a deferred
K-inner sum outside the K loop (ADR-0064 Rev2 D8 single-op FIXED path
exercised across DMA, GEMM, and MATH engines).
"""
from __future__ import annotations

from math import ceil

from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (  # noqa: E501
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _primitive,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import (  # noqa: E501
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel as _tiled,
)
from kernbench.common.pe_commands import (
    DmaReadCmd, GemmCmd, PeCpuOverheadCmd,
)
from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
from kernbench.triton_emu.tl_context import TLContext, run_kernel

_C, _P = 8, 8
_S_KV = 131_072
_D_HEAD, _H_Q, _H_KV, _T_Q = 128, 8, 1, 1
_TILE_S_KV = 1024
_MAC = 16


def _run(kernel, S_kv: int):
    tl = TLContext(
        pe_id=0, num_programs=_P, cost_model=DEFAULT_PE_COST_MODEL,
        cube_id=6, num_cubes=_C, scratch_base=0x200000, scratch_size=1 << 20,
    )
    run_kernel(
        kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
        _T_Q, S_kv, _H_Q, _H_KV, _D_HEAD, _C, _P,
    )
    return tl.commands


def _block_counts():
    """Return (n_tiles, qk_blocks, pv_blocks) for the fixed test params."""
    G = _H_Q // _H_KV
    S_local = _S_KV // (_C * _P)
    n_tiles = (S_local + _TILE_S_KV - 1) // _TILE_S_KV
    tile_s = min(_TILE_S_KV, S_local)
    M = G * _T_Q
    qk = ceil(M / _MAC) * ceil(tile_s / _MAC) * ceil(_D_HEAD / _MAC)
    pv = ceil(M / _MAC) * ceil(_D_HEAD / _MAC) * ceil(tile_s / _MAC)
    return n_tiles, qk, pv


def test_tiled_gemm_count_matches_block_formula():
    """One GemmCmd per 16³ block across both matmuls."""
    n_tiles, qk, pv = _block_counts()
    expected = n_tiles * (qk + pv)

    cmds = _run(_tiled, _S_KV)
    got = sum(1 for c in cmds if isinstance(c, GemmCmd))
    assert got == expected, f"GemmCmd count: got {got}, expected {expected}"


def test_tiled_dma_count_matches_streamed_formula():
    """Streamed operand DMAs — Q·Kᵀ blocks pay 2 DMAs (A+B slice), P·V
    blocks pay 1 DMA (V slice only; the P side is TCM-resident post-softmax).
    Each matmul also carries full-shape coarse loads for the GemmCmd
    operand handles (2 for Q·Kᵀ: Q + K_T, 1 for P·V: V only)."""
    n_tiles, qk, pv = _block_counts()
    expected = n_tiles * (2 * qk + 1 * pv + 3)

    cmds = _run(_tiled, _S_KV)
    got = sum(1 for c in cmds if isinstance(c, DmaReadCmd))
    assert got == expected, f"DmaReadCmd count: got {got}, expected {expected}"


def test_tiled_amplifies_dispatch_vs_primitive():
    """The streamed tiled kernel emits ≫40× more PE_CPU dispatches than
    the coarse primitive — the whole point of adding this variant."""
    n_tiled = sum(1 for c in _run(_tiled, _S_KV)
                  if isinstance(c, PeCpuOverheadCmd))
    n_prim = sum(1 for c in _run(_primitive, _S_KV)
                 if isinstance(c, PeCpuOverheadCmd))
    assert n_tiled > 40 * n_prim, (
        f"expected >40× dispatch amplification, got {n_tiled} vs {n_prim}"
    )