"""
Native CuTeDSL SWA Decode Attention Kernel for DeepSeek-V4 on Blackwell (SM100).

FUSED kernel: paged KV read + Q*K^T + online softmax + V accumulation.
fp8 dequant is done in a batched pre-step on the host side (fast with torch ops).
Future optimization: fuse the fp8 dequant into the kernel using vectorized loads.

CTA mapping: one CTA per (decode_token, q_head_group).
  - 128 Q heads / 16 per group = 8 groups per token
  - Grid: (num_head_groups, num_decode_tokens, 1)
"""

import torch
from typing import Optional

try:
    import cutlass
    import cutlass.cute as cute
    import cutlass.torch as cutlass_torch
    import cutlass.utils as utils
    import cuda.bindings.driver as cuda
    HAS_CUTEDSL = True
except ImportError:
    HAS_CUTEDSL = False

_compiled_kernel_cache = {}

HEAD_GROUP = 16
KV_TILE = 16
HEAD_DIM = 512
NUM_THREADS = 128


def native_swa_decode_attention(
    q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens,
    block_size, scale, window_size=128,
):
    """Native SWA decode attention.

    Pre-dequantizes fp8 KV cache to bf16 in a batched operation,
    then launches the CuTeDSL attention kernel on bf16 data.
    """
    num_tokens, NH, HD = q.shape
    device = q.device

    if not HAS_CUTEDSL:
        return _fallback_batched_sdp(q, swa_kv_cache, swa_inv_scale,
                                     swa_indices, swa_lens, block_size,
                                     scale, window_size)

    q = q.contiguous()
    swa_indices = swa_indices.contiguous()
    swa_lens = swa_lens.contiguous()

    # Pre-dequantize fp8 KV cache to bf16
    # This is a batched gather + dequant: fast on GPU
    if swa_indices.dim() == 3:
        swa_indices_2d = swa_indices.squeeze(0)[:num_tokens]
    else:
        swa_indices_2d = swa_indices[:num_tokens]

    max_len = swa_lens[:num_tokens].max().item()
    if max_len <= 0:
        return torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device)

    # Clamp to window_size
    max_len = min(max_len, window_size)

    # Gather all KV indices: (num_tokens, max_len)
    safe_indices = swa_indices_2d[:, :max_len].clamp(min=0)
    block_indices = safe_indices // block_size
    offsets = safe_indices % block_size

    # Batched gather + dequant
    kv_raw = swa_kv_cache[block_indices, offsets]  # (T, max_len, HD) fp8
    if swa_kv_cache.dtype == torch.uint8:
        kv_raw = kv_raw.view(torch.float8_e4m3fn)
    inv_scales = swa_inv_scale[safe_indices]  # (T, max_len, 1)
    kv_bf16 = (kv_raw.to(torch.bfloat16) * inv_scales).to(torch.bfloat16)

    # Pad to window_size if needed
    if max_len < window_size:
        pad = torch.zeros(num_tokens, window_size - max_len, HD,
                          dtype=torch.bfloat16, device=device)
        kv_bf16 = torch.cat([kv_bf16, pad], dim=1)

    # kv_bf16 is now (num_tokens, window_size, HD) bf16
    output = torch.zeros(num_tokens, NH, HD, dtype=torch.bfloat16, device=device)

    cache_key = (num_tokens, NH, HD, window_size, str(device))

    if cache_key not in _compiled_kernel_cache:
        def to_cute(t):
            ct = cutlass_torch.from_dlpack(t)
            return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))

        q_c = to_cute(q)
        kv_c = to_cute(kv_bf16)
        len_c = to_cute(swa_lens[:num_tokens])
        out_c = to_cute(output)
        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

        scale_tensor = torch.tensor([scale], dtype=torch.float32, device=device)
        scale_c = to_cute(scale_tensor)

        kernel = BlackwellSWADecodeKernel(
            head_dim=HD, head_group=HEAD_GROUP, kv_tile=KV_TILE,
            window_size=window_size,
        )

        compiled = cute.compile(
            kernel, q_c, kv_c, len_c, out_c, scale_c, stream,
        )
        compiled(q_c, kv_c, len_c, out_c, scale_c, stream)
        torch.cuda.synchronize()

        _compiled_kernel_cache[cache_key] = {'compiled': compiled}

    entry = _compiled_kernel_cache[cache_key]
    compiled = entry['compiled']

    def to_cute(t):
        ct = cutlass_torch.from_dlpack(t)
        return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))

    q_c = to_cute(q)
    kv_c = to_cute(kv_bf16)
    len_c = to_cute(swa_lens[:num_tokens])
    out_c = to_cute(output)
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
    scale_tensor = torch.tensor([scale], dtype=torch.float32, device=device)
    scale_c = to_cute(scale_tensor)

    compiled(q_c, kv_c, len_c, out_c, scale_c, stream)
    return output


def _fallback_batched_sdp(
    q, swa_kv_cache, swa_inv_scale, swa_indices, swa_lens,
    block_size, scale, window_size,
):
    num_tokens, NH, HD = q.shape
    device = q.device

    if swa_indices.dim() == 3:
        swa_indices = swa_indices.squeeze(0)

    safe_indices = swa_indices[:num_tokens].clamp(min=0)
    block_indices = safe_indices // block_size
    offsets = safe_indices % block_size

    kv_raw = swa_kv_cache[block_indices, offsets]
    if swa_kv_cache.dtype == torch.uint8:
        kv_raw = kv_raw.view(torch.float8_e4m3fn)
    inv_scales = swa_inv_scale[safe_indices]
    kv_bf16 = (kv_raw.to(torch.bfloat16) * inv_scales).to(torch.bfloat16)

    pos_range = torch.arange(window_size, device=device).unsqueeze(0)
    len_mask = pos_range >= swa_lens[:num_tokens].unsqueeze(1)
    invalid_mask = swa_indices[:num_tokens] < 0
    attn_mask = len_mask | invalid_mask
    float_mask = torch.zeros(attn_mask.shape, dtype=torch.bfloat16, device=device)
    float_mask[attn_mask] = float('-inf')

    q_t = q.permute(1, 0, 2)
    q_batch = q_t.reshape(NH * num_tokens, 1, HD)
    kv_expanded = kv_bf16.unsqueeze(0).expand(NH, num_tokens, window_size, HD)
    k_batch = kv_expanded.reshape(NH * num_tokens, window_size, HD)
    v_batch = k_batch

    mask_batch = float_mask.unsqueeze(0).unsqueeze(2).expand(
        NH, num_tokens, 1, window_size
    ).reshape(NH * num_tokens, 1, window_size)

    out = torch.nn.functional.scaled_dot_product_attention(
        q_batch, k_batch, v_batch,
        attn_mask=mask_batch,
        is_causal=False,
        scale=scale,
    )

    return out.reshape(NH, num_tokens, HD).permute(1, 0, 2)


if HAS_CUTEDSL:

    class BlackwellSWADecodeKernel:
        def __init__(self, head_dim=HEAD_DIM, head_group=HEAD_GROUP,
                     kv_tile=KV_TILE, window_size=128):
            self._head_dim = head_dim
            self._head_group = head_group
            self._kv_tile = kv_tile
            self._window_size = window_size
            self._num_threads = NUM_THREADS

        @cute.jit
        def __call__(self, mQ, mKV, mLens, mO, mScale, stream):
            num_tokens = mQ.shape[0]
            num_head_groups = mQ.shape[1] // self._head_group
            grid_dim = (num_head_groups, num_tokens, 1)

            self._kernel(
                mQ, mKV, mLens, mO, mScale,
            ).launch(
                grid=grid_dim,
                block=[self._num_threads, 1, 1],
                stream=stream,
            )

        @cute.kernel
        def _kernel(
            self,
            mQ: cute.Tensor,       # (T, NH, HD) bf16
            mKV: cute.Tensor,      # (T, WS, HD) bf16 - pre-dequantized
            mLens: cute.Tensor,    # (T,) int64
            mO: cute.Tensor,       # (T, NH, HD) bf16
            mScale: cute.Tensor,   # (1,) f32
        ):
            tidx, _, _ = cute.arch.thread_idx()
            hg_idx, tok_idx, _ = cute.arch.block_idx()

            HG = self._head_group
            HD = self._head_dim
            KT = self._kv_tile
            WS = self._window_size
            softmax_scale = mScale[0]

            # ── Shared memory ──────────────────────────────────────
            @cute.struct
            class SharedStorage:
                kv_tile: cute.struct.MemRange[cutlass.BFloat16, KT * HD]

            smem = utils.SmemAllocator()
            storage = smem.allocate(SharedStorage)

            sKV = cute.make_tensor(
                storage.kv_tile.data_ptr(),
                cute.make_layout((KT, HD), stride=(HD, 1)),
            )

            # ── Read valid KV length ───────────────────────────────
            swa_len = mLens[tok_idx]
            has_kv = swa_len > 0

            # ── Load Q into registers: (HG, HD) ───────────────────
            q_reg = cute.make_rmem_tensor((HG, HD), cutlass.BFloat16)
            for h in cutlass.range_constexpr(HG):
                qh = hg_idx * HG + h
                for d in range(HD):
                    q_reg[h, d] = mQ[tok_idx, qh, d]

            # ── Output accumulator: (HG, HD) f32 ──────────────────
            acc_O = cute.make_rmem_tensor((HG, HD), cutlass.Float32)
            acc_O.fill(0.0)

            # ── Online softmax state: (HG,) f32 ───────────────────
            row_max = cute.make_rmem_tensor((HG,), cutlass.Float32)
            row_sum = cute.make_rmem_tensor((HG,), cutlass.Float32)
            row_max.fill(-1e30)
            row_sum.fill(0.0)

            # ── Stream KV tiles ────────────────────────────────────
            max_tiles = (WS + KT - 1) // KT

            for tile_idx in range(max_tiles):
                tile_start = tile_idx * KT

                # Load bf16 KV from contiguous tensor to smem
                for kv_pos in range(KT):
                    global_kv = tile_start + kv_pos
                    for d in range(HD):
                        valid = global_kv < swa_len
                        val = cutlass.BFloat16(0.0)
                        if valid:
                            val = mKV[tok_idx, global_kv, d]
                        sKV[kv_pos, d] = val

                cute.arch.sync_threads()

                # Q * K^T: (HG, KT) scores
                scores = cute.make_rmem_tensor((HG, KT), cutlass.Float32)
                scores.fill(0.0)

                for h in cutlass.range_constexpr(HG):
                    for kv_pos in range(KT):
                        dot = cutlass.Float32(0.0)
                        for d in range(HD):
                            q_val = q_reg[h, d].to(cutlass.Float32)
                            k_val = sKV[kv_pos, d].to(cutlass.Float32)
                            dot = dot + q_val * k_val
                        scores[h, kv_pos] = dot * softmax_scale

                # Online softmax update
                for h in cutlass.range_constexpr(HG):
                    tile_max = cutlass.Float32(-1e30)
                    for kv_pos in range(KT):
                        s = scores[h, kv_pos]
                        if s > tile_max:
                            tile_max = s

                    new_max = row_max[h]
                    if tile_max > new_max:
                        new_max = tile_max

                    rescale = cutlass.Float32(0.0)
                    if row_max[h] > cutlass.Float32(-1e29):
                        rescale = cute.exp(row_max[h] - new_max)

                    for d in range(HD):
                        acc_O[h, d] = acc_O[h, d] * rescale
                    row_sum[h] = row_sum[h] * rescale

                    for kv_pos in range(KT):
                        exp_score = cute.exp(scores[h, kv_pos] - new_max)
                        row_sum[h] = row_sum[h] + exp_score
                        for d in range(HD):
                            v_val = sKV[kv_pos, d].to(cutlass.Float32)
                            acc_O[h, d] = acc_O[h, d] + exp_score * v_val

                    row_max[h] = new_max

                cute.arch.sync_threads()

            # ── Normalize and write output ─────────────────────────
            for h in cutlass.range_constexpr(HG):
                qh = hg_idx * HG + h
                for d in range(HD):
                    val_f32 = cutlass.Float32(0.0)
                    if has_kv and row_sum[h] > cutlass.Float32(1e-30):
                        val_f32 = acc_O[h, d] / row_sum[h]
                    mO[tok_idx, qh, d] = val_f32.to(cutlass.BFloat16)