"""GQA decode kernel — Case 6, **primitive hand-tiled 16×16×16** (per-block DMA + MAC-side accumulate).

Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is
that each local-attention matmul is *hand-blocked into 16×16×16 GEMMs*
(mac=16), **with each block re-fetching its operand slices from HBM**
(no operand cache), and each (mi, ni) output tile updated **implicitly**
by successive K-inner ``GemmCmd`` writes into the shared (M, N) output —
mirroring a MAC-array accumulator register that latches across K-inner
cycles on real accelerators.

Per K-inner block: two ``tl.load`` DMAs (Q and K slice for Q·Kᵀ, or one
V slice for P·V) → one ``GemmCmd`` emitted directly via ``tl._emit`` with
``out`` bound to the shared (M, N) handle and ``m/k/n`` overridden to
the 16³ block dims. All K-inner iterations at a given (mi, ni) write
into the same output tile; no compiler-emitted ``MathCmd`` accumulator
chain.

This models a strict-streaming architecture (no HBM-operand cache) — the
worst-case dispatch endpoint: every 16³ compute block pays the ADR-0064
D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM,
exposing the full PE_CPU dispatch pressure that composite forms
(ADR-0065) absorb into PE_SCHEDULER.

FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute
of its block; the blocks sum to the full matmul); end-to-end compute
time is unchanged vs the coarse primitive — only the PE_CPU command
count and its dispatch cycles grow. Inputs are zero (decode bench
convention), so the "overwrite instead of accumulate" is semantically
equivalent to sum (0 = 0 + 0), and ``GemmCmd`` writes populate the
shared ``out`` handle so the downstream softmax's strict ``MathCmd``
reads succeed — the kernel runs in engine (data) mode.
"""
from __future__ import annotations

from math import ceil

from kernbench.common.pe_commands import GemmCmd
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
    _merge_running,
    reduce_mlo,
    root_cube_for,
)

TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
MAC = 16          # 16×16×16 MAC-array blocking granularity.


def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl):
    """Q·Kᵀ hand-blocked, per-block DMA of both operands, MAC-side accumulate.

    For each 16³ block (mi, ni, ki): load 16×16 A and B slices from HBM,
    emit a ``GemmCmd`` directly with ``out`` bound to the shared (M, N)
    handle and ``m/k/n`` overridden to the 16³ block dims. All K-inner
    iterations at a given (mi, ni) write into the same shared ``out`` tile
    — accumulation is implicit (MAC-side latching, mirroring how real
    accelerators feed a per-tile accumulator register instead of running
    a compiler-emitted MathCmd add chain).

    Under the zero-input decode-bench convention, the "overwrite instead
    of accumulate" is semantically equivalent to sum (0 = 0 + 0), and
    ``GemmCmd`` writes populate ``out.addr`` in ``MemoryStore`` so the
    downstream softmax's strict ``MathCmd`` reads succeed. Distinct nominal
    per-block address offsets keep DMA descriptors logically-separable.
    """
    # Full-shape operand handles for data-mode correctness. Under the
    # simulator's DataExecutor, GemmCmd computes np.matmul over these
    # shapes and writes an (M, N) result to out.addr — populating the
    # shared handle with a valid region so the downstream softmax's
    # strict MathCmd reads succeed. In engine timing, the m/k/n
    # override on each GemmCmd charges only the 16³ block work.
    A_full = tl.load(a_ptr, shape=(M, K), dtype="f16")
    B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
    out = tl._make_compute_out(shape=(M, N), dtype="f16")
    n_m = ceil(M / MAC)
    n_k = ceil(K / MAC)
    n_n = ceil(N / MAC)
    for mi in range(n_m):
        bm = min(MAC, M - mi * MAC)
        for ni in range(n_n):
            bn = min(MAC, N - ni * MAC)
            # Recycle per (mi, ni) — the per-block slice DMAs live only
            # during this tile's processing; the shared out survives.
            with tl.scratch_scope():
                for ki in range(n_k):
                    bk = min(MAC, K - ki * MAC)
                    # Per-block DMAs model the streaming architecture
                    # (no HBM-operand cache). Row-major byte offsets:
                    # A is (M, K), B is (K, N). Nominal handles — the
                    # GemmCmd below uses the full-shape handles above
                    # for data-mode compute; these per-block loads pay
                    # their ADR-0064 dispatch cost as descriptor work.
                    _ = tl.load(
                        a_ptr + mi * MAC * K * 2 + ki * MAC * 2,
                        shape=(bm, bk), dtype="f16",
                    )
                    _ = tl.load(
                        b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
                        shape=(bk, bn), dtype="f16",
                    )
                    tl._emit(GemmCmd(
                        a=A_full, b=B_full, out=out,
                        m=bm, k=bk, n=bn,
                    ))
    return out


def _blocked_dot_pv_streamed(A_handle, b_ptr, M, K, N, *, tl):
    """P·V hand-blocked; only V streams (P is TCM-resident post-softmax).

    Same structure as the Q·Kᵀ variant — per block: 1 DMA (V slice) + 1
    ``GemmCmd`` writing to the shared (M, N) ``out`` handle with implicit
    K-inner accumulation via successive block writes.

    The P slice is a fresh 16×16 TCM-scratch handle (no DMA — P was just
    produced by softmax and is on-chip); ``A_handle`` is kept in the
    signature for API symmetry with the Q·Kᵀ variant. Zero-input
    convention applies throughout.
    """
    # A_handle (= exp_scores from softmax) is TCM-resident with
    # shape (M, K), already populated by tl.exp's DataExecutor. Only V
    # needs a full-shape coarse load from HBM for data-mode correctness.
    B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
    out = tl._make_compute_out(shape=(M, N), dtype="f16")
    n_m = ceil(M / MAC)
    n_k = ceil(K / MAC)
    n_n = ceil(N / MAC)
    for _mi in range(n_m):
        for ni in range(n_n):
            bn = min(MAC, N - ni * MAC)
            with tl.scratch_scope():
                for ki in range(n_k):
                    bk = min(MAC, K - ki * MAC)
                    # Per-block V load — streaming-architecture dispatch
                    # cost. GemmCmd below uses B_full for compute.
                    _ = tl.load(
                        b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
                        shape=(bk, bn), dtype="f16",
                    )
                    bm = min(MAC, M - _mi * MAC)
                    tl._emit(GemmCmd(
                        a=A_handle, b=B_full, out=out,
                        m=bm, k=bk, n=bn,
                    ))
    return out


def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel(
    q_ptr: int,
    k_ptr: int,
    v_ptr: int,
    o_ptr: int,
    T_q: int,
    S_kv: int,
    h_q: int,
    h_kv: int,
    d_head: int,
    C: int,
    P: int,
    *,
    tl,
) -> None:
    """Case-6 decode, primitive-TILED streamed (per-block DMA + 16³ GEMM)."""
    G = h_q // h_kv
    n_ranks = C * P
    S_local = S_kv // n_ranks
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)

    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
    KV_ROW_BYTES = d_head * 2  # f16

    # Bootstrap tile (tile 0). Establishes persistent (m_local, l_local,
    # O_local). Cannot fold into Tiles 1..N loop: persistent tensors must
    # live outside tl.scratch_scope or scope teardown discards them.
    tile_s0 = min(TILE_S_KV, S_local)
    scores = _blocked_dot_qk_streamed(
        q_ptr, k_ptr, G * T_q, d_head, tile_s0, tl=tl,
    )
    m_local = tl.max(scores, axis=-1)
    centered = scores - m_local
    exp_scores = tl.exp(centered)
    l_local = tl.sum(exp_scores, axis=-1)
    O_local = _blocked_dot_pv_streamed(
        exp_scores, v_ptr, G * T_q, tile_s0, d_head, tl=tl,
    )

    for tile_idx in range(1, n_tiles):
        tile_start = tile_idx * TILE_S_KV
        tile_s = min(TILE_S_KV, S_local - tile_start)
        with tl.scratch_scope():
            scores_t = _blocked_dot_qk_streamed(
                q_ptr, k_ptr + tile_start * KV_ROW_BYTES,
                G * T_q, d_head, tile_s, tl=tl,
            )
            m_tile = tl.max(scores_t, axis=-1)
            centered_t = scores_t - m_tile
            exp_scores_t = tl.exp(centered_t)
            l_tile = tl.sum(exp_scores_t, axis=-1)
            O_tile = _blocked_dot_pv_streamed(
                exp_scores_t, v_ptr + tile_start * KV_ROW_BYTES,
                G * T_q, tile_s, d_head, tl=tl,
            )
            m_new, l_new, O_new = _merge_running(
                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
            )
            tl.copy_to(m_local, m_new)
            tl.copy_to(l_local, l_new)
            tl.copy_to(O_local, O_new)

    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)

    if pe_id == 0 and cube_id == root_cube_for(C):
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
