"""GQA decode kernel: short context, attention only, multi-tile per PE (1).

Unified A1/A2/A4/B decode mapping per ADR-0070 (phase mirror of
``_gqa_attention_prefill_short.py``). Mode selected at launch via
``kv_per_cube ∈ {1, 2, 4, 8}``:

    Mode  kv_per_cube  C        group_size  Reduce topology
    ----  -----------  -------  ----------  -----------------------------
    A1            1    h_kv      P (=8)     row 0 + col bridge + row 1
    A2            2    h_kv/2    P/2 (=4)   row chain only
    A4            4    h_kv/4    P/4 (=2)   single intra_W hop
    B             8    1         1          (no reduce, single PE)

Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
PEs; each PE owns ``S_local = S_kv/group_size`` KV tokens (sequence
shard). FA2 fuses ``G = h_q/h_kv`` Q heads into the M dim of one GEMM
per tile. After ``n_tiles_per_pe = S_local/TILE_S_KV`` tiles, partials
``(m, ℓ, O)`` chain-reduce up to group root (PE 0), which normalizes
and stores the cube's O slab.

Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
deploy places shards in HBM but the kernel must address them.

Tensor layouts (host-side, mode-invariant byte totals):
  - Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``, T_q=1.
    dp=(cube=column_wise, pe=replicate). Caller pre-scales by
    ``1/sqrt(d_head)``.
  - K: ``(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)`` tile-major.
    dp=(cube=row_wise, pe=row_wise).
  - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
  - O: same dp as Q; only group root (pe_in_group==0) stores.

Configuration constraints:
  kv_per_cube ∈ {1, 2, 4, 8},  T_q == 1,  P == 8,
  C == h_kv/kv_per_cube,  h_q % h_kv == 0,
  S_kv % (group_size·TILE_S_KV) == 0.

Out of scope: causal mask (decode is auto-causal), f32 accumulator.
"""
from __future__ import annotations


TILE_S_KV = 1024


def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
                     h_q: int, h_kv: int, S_kv: int) -> None:
    """Validate decode kernel configuration before the run.
    """
    if kv_per_cube not in (1, 2, 4, 8):
        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
    if T_q != 1:
        raise ValueError(f"decode requires T_q == 1; got {T_q}")
    if P != 8:
        raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
    if h_q % h_kv != 0:
        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
    if C != h_kv // kv_per_cube:
        raise ValueError(
            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
        )
    group_size = P // kv_per_cube
    if S_kv % (group_size * TILE_S_KV) != 0:
        raise ValueError(
            f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
            f"({group_size * TILE_S_KV}); each PE's sequence shard must "
            f"be a whole number of tiles"
        )


def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
    m_new = tl.maximum(m_local, m_other)
    scale_old = tl.exp(m_local - m_new)
    scale_new = tl.exp(m_other - m_new)
    l_new = l_local * scale_old + l_other * scale_new
    O_new = O_local * scale_old + O_other * scale_new
    return m_new, l_new, O_new


def gqa_attention_decode_short_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,
    kv_per_cube: int,
    cube_base: int = 0,
    *,
    tl,
) -> None:
    """Unified decode: sequence-shard + chain reduce + FA2 (ADR-0070).
    """
    group_size = P // kv_per_cube
    S_local = S_kv // group_size
    n_tiles_per_pe = S_local // TILE_S_KV

    G = h_q // h_kv
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    # User-local cube index: a batched user placed at cube_start=cube_base
    # addresses its own shards from a 0-based cube index (default 0 = single user).
    cube_local = cube_id - cube_base
    pe_in_group = pe_id % group_size
    group_id_in_cube = pe_id // group_size

    KV_ROW_BYTES = d_head * 2
    K_TILE_BYTES = d_head * TILE_S_KV * 2
    Q_ROW_BYTES = G * d_head * 2
    K_HEAD_BYTES = S_kv * d_head * 2
    V_HEAD_BYTES = S_kv * KV_ROW_BYTES

    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
    q_base = q_ptr + cube_local * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
    k_shard_base = (k_ptr
                    + cube_local * kv_per_cube * K_HEAD_BYTES
                    + group_id_in_cube * K_HEAD_BYTES
                    + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
    v_shard_base = (v_ptr
                    + cube_local * kv_per_cube * V_HEAD_BYTES
                    + group_id_in_cube * V_HEAD_BYTES
                    + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)

    Q = tl.load(q_base, shape=(G, d_head), dtype="f16")

    # ── Tile 0: establish (m, ℓ, O) ──────────────────────────────────
    K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
    V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
    scores = tl.dot(Q, K_T)
    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 = tl.dot(exp_scores, V)

    # ── Tiles 1..n_tiles_per_pe-1: sweep this PE's sequence shard ──
    for tile_idx in range(1, n_tiles_per_pe):
        with tl.scratch_scope():
            K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
                            shape=(d_head, TILE_S_KV), dtype="f16")
            V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
                          shape=(TILE_S_KV, d_head), dtype="f16")
            scores_t = tl.dot(Q, K_T_t)
            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 = tl.dot(exp_scores_t, V_t)
            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)

    # ── 2x4 mesh chain reduce geometry within the group ──
    group_cols = min(4, group_size)
    group_rows = (group_size + group_cols - 1) // group_cols
    pe_col_in_group = pe_in_group % group_cols
    pe_row_in_group = pe_in_group // group_cols

    # Row chain (intra_W, leftward) — every group with group_cols > 1.
    if group_cols > 1:
        if pe_col_in_group < group_cols - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_col_in_group > 0:
            tl.send(dir="intra_W", src=m_local)
            tl.send(dir="intra_W", src=l_local)
            tl.send(dir="intra_W", src=O_local)

    # Col bridge (intra_N, row-1 col-0 → row-0 col-0) — only A1 (group_rows > 1).
    if pe_col_in_group == 0 and group_rows > 1:
        if pe_row_in_group < group_rows - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_row_in_group > 0:
            tl.send(dir="intra_N", src=m_local)
            tl.send(dir="intra_N", src=l_local)
            tl.send(dir="intra_N", src=O_local)

    # ── Final normalize + store (group root, pe_in_group == 0) ──
    if pe_in_group == 0:
        O_final = O_local / l_local
        o_base = (o_ptr
                  + cube_local * kv_per_cube * Q_ROW_BYTES
                  + group_id_in_cube * Q_ROW_BYTES)
        tl.store(o_base, O_final)
