"""Request-owned, raw-bit-preserving summary-indexer decode materialization."""
import torch
import triton
import triton.language as tl


@triton.jit
def _gather(K, V, Blocks, OK, OV, Mask, Candidates, Safe,
            L, H: tl.constexpr, D: tl.constexpr,
            N: tl.constexpr, W: tl.constexpr, R: tl.constexpr,
            START, KS0: tl.constexpr, KS1: tl.constexpr,
            KS2: tl.constexpr, KS3: tl.constexpr,
            VS0: tl.constexpr, VS1: tl.constexpr, VS2: tl.constexpr,
            VS3: tl.constexpr, BS0: tl.constexpr, BS2: tl.constexpr,
            TILE: tl.constexpr, DIM: tl.constexpr):
    bh = tl.program_id(0)
    b, h = bh // H, bh % H
    t = tl.program_id(1) * TILE + tl.arange(0, TILE)
    d = tl.arange(0, DIM)
    remote = t < N * R
    block = tl.load(Blocks + b * BS0 + (t // R) * BS2,
                    mask=remote & (t < W), other=-1).to(tl.int64)
    index = tl.where(remote, block * R + t % R, START + t - N * R)
    valid = tl.where(remote, (block >= 0) & (index < L),
                     (index >= 0) & (index < L))
    candidate = tl.where(valid, index, L)
    safe = tl.minimum(candidate, L - 1)
    k = tl.load(K + b * KS0 + h * KS1 + safe[:, None] * KS2 + d[None, :] * KS3,
                mask=(t[:, None] < W) & (d[None, :] < D), other=0)
    v = tl.load(V + b * VS0 + h * VS1 + safe[:, None] * VS2 + d[None, :] * VS3,
                mask=(t[:, None] < W) & (d[None, :] < D), other=0)
    out = (bh * W + t[:, None]) * D + d[None, :]
    tl.store(OK + out, k, mask=(t[:, None] < W) & (d[None, :] < D))
    tl.store(OV + out, v, mask=(t[:, None] < W) & (d[None, :] < D))
    if h == 0:
        tl.store(Mask + b * W + t, valid, mask=t < W)
        tl.store(Candidates + b * W + t, candidate, mask=t < W)
        tl.store(Safe + b * W + t, safe, mask=t < W)


def fused_summary_decode_gather(key, value, selected, *, raw_block_size,
                                local_window_tokens, local_alignment_tokens,
                                workspace):
    """Write every output slot, including masked slots, on the input device.

    Masked K/V equal the reference's last-live-token copy. Exact compact shapes
    and contiguous strides preserve SDPA dispatch. Storage is cache-local;
    shape/dtype/device changes replace it. Returned views live until next call.
    """
    if key.shape != value.shape or key.ndim != 4:
        raise ValueError('decode K/V must share [B,H,L,D]')
    b, h, length, dim = key.shape
    if selected.ndim != 3 or selected.shape[:2] != (b, 1):
        raise ValueError('decode selected blocks must have [B,1,K] layout')
    alignment = raw_block_size if local_alignment_tokens is None else local_alignment_tokens
    if raw_block_size <= 0 or local_window_tokens <= 0 or alignment <= 0 or alignment % raw_block_size:
        raise ValueError('invalid block/window/alignment')
    if length == 0:
        raise ValueError('decode requires live K/V')
    n = selected.shape[-1]
    width = n * raw_block_size + local_window_tokens + alignment - 1
    signature = (b, h, width, dim, key.dtype, key.device)
    if workspace.get('signature') != signature:
        workspace.clear()
        workspace.update(signature=signature,
                         key=torch.empty((b,h,width,dim),dtype=key.dtype,device=key.device),
                         value=torch.empty((b,h,width,dim),dtype=value.dtype,device=value.device),
                         mask=torch.empty((b,1,1,width),dtype=torch.bool,device=key.device),
                         candidates=torch.empty((b,width),dtype=torch.int64,device=key.device),
                         safe=torch.empty((b,width),dtype=torch.int64,device=key.device))
    start = max(length-local_window_tokens,0) // alignment * alignment
    with torch.cuda.device(key.device):
        _gather[(b*h,triton.cdiv(width,16))](
            key,value,selected,workspace['key'],workspace['value'],workspace['mask'],
            workspace['candidates'],workspace['safe'],length,h,dim,n,width,raw_block_size,start,
            *key.stride(),*value.stride(),selected.stride(0),selected.stride(2),
            16,triton.next_power_of_2(dim),num_warps=4)
    return workspace['key'],workspace['value'],workspace['mask']