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

Unified A1/A2/A4/B prefill mapping per ADR-0070 (supersedes
ADR-0060 §B.split.2 prefill-short clause). Mode selected at launch
via ``kv_per_cube ∈ {1, 2, 4, 8}``:

    Mode  kv_per_cube  C        group_size  Broadcast 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_E hop
    B             8    1         1          (no broadcast, 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; the group's T_q rows are floor-balanced across its PEs (Q-tile
split). FA2 fuses ``G = h_q/h_kv`` Q heads into one batched GEMM per
PE per tile. Group root (pe_in_group==0) loads K/V from HBM and IPCQ-
broadcasts each tile to the rest of its group; q-tiles are
independent so no intra-group reduce — every PE writes its own
``(T_q_pe·G, d_head)`` slab to disjoint rows of O.

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)``.
    dp=(cube=column_wise, pe=replicate). Caller pre-scales by
    ``1/sqrt(d_head)``.
  - K: ``(h_kv·n_tiles·d_head, TILE_S_KV)`` tile-major.
    dp=(cube=row_wise, pe=replicate).
  - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
  - O: same dp as Q; every PE stores.

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

Out of scope: causal mask, f32 accumulator. See ADR-0070 §Known
Limitations.
"""
from __future__ import annotations


TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).


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 prefill 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 P != 8:
        raise ValueError(f"2x4 mesh broadcast 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 T_q < group_size:
        raise ValueError(
            f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
            f"has q-tile work for the broadcast chain"
        )
    if S_kv % TILE_S_KV != 0:
        raise ValueError(
            f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
        )


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_prefill_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,
    *,
    tl,
) -> None:
    """Unified prefill: Q-tile split + FA2 + IPCQ KV broadcast (ADR-0070)."""
    group_size = P // kv_per_cube

    G = h_q // h_kv
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    pe_in_group = pe_id % group_size
    group_id_in_cube = pe_id // group_size

    # Floor-balanced Q-tile partition.
    q_start = (T_q * pe_in_group) // group_size
    q_end = (T_q * (pe_in_group + 1)) // group_size
    T_q_pe = q_end - q_start

    n_tiles = S_kv // TILE_S_KV

    Q_ROW_BYTES = G * d_head * 2
    KV_ROW_BYTES = d_head * 2
    K_TILE_BYTES = d_head * TILE_S_KV * 2
    K_HEAD_BYTES = n_tiles * K_TILE_BYTES
    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_id * kv_per_cube * T_q * Q_ROW_BYTES
              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
    k_head_shard_base = (k_ptr
                         + cube_id * kv_per_cube * K_HEAD_BYTES
                         + group_id_in_cube * K_HEAD_BYTES)
    v_head_shard_base = (v_ptr
                         + cube_id * kv_per_cube * V_HEAD_BYTES
                         + group_id_in_cube * V_HEAD_BYTES)

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

    # 2x4 mesh broadcast 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

    # ──────────────────────────────────────────────────────────
    # Tile 0: KV broadcast + establish persistent (m, ℓ, O).
    # Persistent state lives OUTSIDE scratch_scope (ADR-0063 §A.3).
    # ──────────────────────────────────────────────────────────
    if pe_in_group == 0:
        K_T = tl.load(k_head_shard_base,
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.load(v_head_shard_base,
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if group_cols > 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
        if group_rows > 1:
            tl.send(dir="intra_S", src=K_T)
            tl.send(dir="intra_S", src=V)
    elif pe_col_in_group == 0 and pe_row_in_group > 0:
        K_T = tl.recv(dir="intra_N",
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.recv(dir="intra_N",
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if group_cols > 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
    else:
        K_T = tl.recv(dir="intra_W",
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.recv(dir="intra_W",
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if pe_col_in_group < group_cols - 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)

    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-1: broadcast + fold into running state.
    # ──────────────────────────────────────────────────────────
    for tile_idx in range(1, n_tiles):
        with tl.scratch_scope():
            if pe_in_group == 0:
                K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if group_cols > 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
                if group_rows > 1:
                    tl.send(dir="intra_S", src=K_T_t)
                    tl.send(dir="intra_S", src=V_t)
            elif pe_col_in_group == 0 and pe_row_in_group > 0:
                K_T_t = tl.recv(dir="intra_N",
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.recv(dir="intra_N",
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if group_cols > 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
            else:
                K_T_t = tl.recv(dir="intra_W",
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.recv(dir="intra_W",
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if pe_col_in_group < group_cols - 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
            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)

    # No intra-group reduce — q-tiles independent. Each PE writes its own
    # slab to disjoint rows of the cube's column-wise O slab.
    O_final = O_local / l_local
    o_base = (o_ptr
              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
    tl.store(o_base, O_final)
