"""Runtime compatibility patch for GLM-5.3 NoPE sparse MLA on SM120/SM121.

The fp8_ds_mla physical cache layout reserves 64 BF16 dimensions for RoPE.
GLM-5.3 uses native NoPE (qk_rope_head_dim=0), so vLLM otherwise forwards an
empty k_pe tensor to concat_and_cache_mla, whose packed-layout kernel requires
pe_dim=64. Pad only the physical cache write; the model and attention kernel
continue to use the architecture-correct logical RoPE dimension of zero.
"""

from __future__ import annotations

import os

import torch
import flashinfer.mla._sparse_mla_sm120 as _sparse_mla_sm120

from vllm.v1.attention.backend import MLAAttentionImpl
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
    _get_workspace_buffer,
)
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
    FlashInferMLASparseSM120Impl,
)
from vllm.v1.attention.backends.mla.sparse_utils import (
    triton_convert_req_index_to_global_index,
)


_original_decode_dsv3_2_dispatchable = (
    _sparse_mla_sm120._decode_dsv3_2_dispatchable
)
_debug = os.getenv("GLM53_GB10_PATCH_DEBUG") == "1"
_debugged_forward = False

if (64, 2176) not in _sparse_mla_sm120._DECODE_DSV3_2_DISPATCH:
    raise RuntimeError(
        "This mock requires the GLM-5.3 H=64/top-k=2176 FlashInfer decode "
        "specialization; use the companion GB10 adapter image."
    )


def _glm53_decode_dsv3_2_dispatchable(
    num_tokens: int,
    num_heads: int,
    topk: int,
    d_qk: int,
    page_block_size: int,
) -> bool:
    # The CLI block size is 256 for kpool storage alignment; vLLM exposes the
    # compressed sparse-MLA cache to FlashInfer as physical 64-token pages.
    result = _original_decode_dsv3_2_dispatchable(
        num_tokens,
        num_heads,
        topk,
        d_qk,
        page_block_size,
    )
    if _debug:
        print(
            "GLM53_PATCH dispatch "
            f"tokens={num_tokens} heads={num_heads} topk={topk} "
            f"d_qk={d_qk} physical_page={page_block_size} -> {result}",
            flush=True,
        )
    return result


_sparse_mla_sm120._decode_dsv3_2_dispatchable = (
    _glm53_decode_dsv3_2_dispatchable
)


_original_do_kv_cache_update = MLAAttentionImpl.do_kv_cache_update


def _glm53_nope_do_kv_cache_update(
    self,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kv_cache_dtype: str,
    k_scale: torch.Tensor,
) -> None:
    if kv_cache_dtype == "fp8_ds_mla" and k_pe.shape[-1] == 0:
        k_pe = k_pe.new_zeros((*k_pe.shape[:-1], 64))
    return _original_do_kv_cache_update(
        self,
        kv_c_normed,
        k_pe,
        kv_cache,
        slot_mapping,
        kv_cache_dtype,
        k_scale,
    )


MLAAttentionImpl.do_kv_cache_update = _glm53_nope_do_kv_cache_update


def _glm53_nope_forward_mqa(
    self,
    q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    kv_c_and_k_pe_cache: torch.Tensor,
    attn_metadata,
    layer,
) -> tuple[torch.Tensor, None]:
    """Run logical NoPE through the physical 576-wide GLM_NSA kernel.

    SM120's GLM_NSA sparse kernel supports the architecture-required top-k
    2048, but its packed cache ABI is 512 latent FP8 values plus 64 BF16 RoPE
    values. Appending zero RoPE values preserves the exact NoPE dot product.
    """
    global _debugged_forward

    if isinstance(q, tuple):
        q = torch.cat(q, dim=-1)

    if self.qk_rope_head_dim != 0:
        raise RuntimeError("GLM-5.3 NoPE patch received a non-NoPE MLA layer")

    q = torch.nn.functional.pad(q, (0, 64))
    num_actual_toks = q.shape[0]
    assert self.topk_indices_buffer is not None
    topk_indices = self.topk_indices_buffer[:num_actual_toks]
    if attn_metadata.topk_tokens != 2048:
        raise RuntimeError(
            f"expected architecture index_topk=2048, got {attn_metadata.topk_tokens}"
        )
    if topk_indices.shape[1] != 2176:
        raise RuntimeError(
            "expected the GLM kpool/tail/alignment buffer width 2176, got "
            f"{topk_indices.shape[1]}"
        )

    converted = triton_convert_req_index_to_global_index(
        attn_metadata.req_id_per_token[:num_actual_toks],
        attn_metadata.block_table,
        topk_indices,
        BLOCK_SIZE=attn_metadata.block_size,
        NUM_TOPK_TOKENS=topk_indices.shape[1],
        return_valid_counts=True,
    )
    topk_indices_physical, valid_counts = converted
    empty_rows = valid_counts == 0
    topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
        empty_rows, 0
    )
    active_lengths = valid_counts.clamp(min=1)
    sparse_topk_capacity = topk_indices_physical.shape[1]
    if _debug and not _debugged_forward:
        print(
            "GLM53_PATCH forward "
            f"q={tuple(q.shape)} topk={tuple(topk_indices.shape)} "
            f"physical_topk={tuple(topk_indices_physical.shape)} "
            f"valid_min={int(valid_counts.min())} "
            f"valid_max={int(valid_counts.max())} "
            f"block_size={attn_metadata.block_size} "
            f"cache={tuple(kv_c_and_k_pe_cache.shape)}",
            flush=True,
        )
        _debugged_forward = True

    output = q.new_empty(
        (num_actual_toks, self.num_heads, self.kv_lora_rank),
        dtype=q.dtype,
    )
    if self._workspace_buffer is None:
        self._workspace_buffer = _get_workspace_buffer(q.device)

    from vllm.utils.flashinfer import (
        flashinfer_trtllm_batch_decode_with_kv_cache_mla,
    )

    out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
        query=q.unsqueeze(1),
        kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
        workspace_buffer=self._workspace_buffer,
        qk_nope_head_dim=self.qk_nope_head_dim,
        kv_lora_rank=self.kv_lora_rank,
        qk_rope_head_dim=64,
        block_tables=topk_indices_physical.unsqueeze(1),
        seq_lens=active_lengths,
        max_seq_len=sparse_topk_capacity,
        out=output.unsqueeze(1),
        bmm1_scale=self.scale,
        bmm2_scale=1.0,
        sparse_mla_top_k=sparse_topk_capacity,
        kv_scale_format=self.kv_scale_format,
    ).squeeze(1)
    out = out.masked_fill(empty_rows[:, None, None], 0)
    return out, None


FlashInferMLASparseSM120Impl.forward_mqa = _glm53_nope_forward_mqa