#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# from collections.abc import Iterable
# mypy: ignore-errors


import torch
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
from vllm.v1.attention.backend import AttentionMetadata  # type: ignore
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
from vllm.v1.attention.backends.utils import PAD_SLOT_ID

from vllm_ascend._310p.ops.fla.chunk_gated_delta_rule import chunk_gated_delta_rule_310
from vllm_ascend._310p.ops.fla.fused_gdn_gating import fused_gdn_gating_pytorch
from vllm_ascend._310p.ops.fla.l2norm import l2norm_310p
from vllm_ascend.ascend_forward_context import _EXTRA_CTX
from vllm_ascend.attention.utils import maybe_save_kv_layer_to_connector
from vllm_ascend.utils import enable_sp


def _zero_padded_tokens(
    tensor: torch.Tensor,
    valid_tokens: torch.Tensor,
    token_dim: int,
) -> torch.Tensor:
    if tensor.numel() == 0:
        return tensor

    token_count = tensor.shape[token_dim]
    if token_count == 0:
        return tensor

    positions = torch.arange(
        token_count,
        device=tensor.device,
        dtype=valid_tokens.dtype,
    )
    valid_mask = positions < valid_tokens.to(device=tensor.device)
    mask_shape = [1] * tensor.ndim
    mask_shape[token_dim] = token_count
    return tensor * valid_mask.reshape(mask_shape).to(dtype=tensor.dtype)


def _flatten_state_indices(
    ssm_state_indices: torch.Tensor,
    cu_seqlens: torch.Tensor,
    total_tokens: int,
) -> torch.Tensor:
    if ssm_state_indices.ndim == 1:
        return ssm_state_indices[:total_tokens].to(torch.int32).contiguous()

    num_seqs = (cu_seqlens[1:] - cu_seqlens[:-1]).shape[0]
    seq_lens = cu_seqlens[1 : num_seqs + 1] - cu_seqlens[:num_seqs]
    ssm_state_indices = ssm_state_indices[:num_seqs]

    # Uniform spec-decode ACL graph uses fixed q_len per request; reshape avoids
    # NPU masked_select which breaks stream capture (aclnnMaskedSelect / 107027).
    if _EXTRA_CTX.capturing or (seq_lens.numel() > 0 and torch.all(seq_lens == seq_lens[0])):
        q_per_seq = ssm_state_indices.shape[1]
        flat = ssm_state_indices[:, :q_per_seq].reshape(-1)
        return flat[:total_tokens].to(torch.int32).contiguous()

    # Eager mixed batches with variable seq_lens: compact on CPU, copy back async.
    ssm_cpu = ssm_state_indices.cpu()
    seq_lens_cpu = seq_lens.cpu()
    q_per_seq = ssm_cpu.shape[1]
    positions = torch.arange(q_per_seq)
    valid = positions.unsqueeze(0) < seq_lens_cpu.unsqueeze(1)
    flat_cpu = ssm_cpu.masked_select(valid).to(torch.int32).contiguous()[:total_tokens]
    if not flat_cpu.is_pinned:
        flat_cpu = flat_cpu.pin_memory()
    flat_dev = torch.empty(flat_cpu.numel(), dtype=torch.int32, device=ssm_state_indices.device)
    flat_dev.copy_(flat_cpu, non_blocking=True)
    return flat_dev.contiguous()


def _mask_padded_recurrent_accepted_tokens(
    num_accepted_tokens: torch.Tensor,
    actual_seq_lengths: torch.Tensor,
) -> torch.Tensor:
    accepted_tokens = num_accepted_tokens[: actual_seq_lengths.shape[0]].to(torch.int32).contiguous()
    return torch.where(
        actual_seq_lengths > 0,
        accepted_tokens,
        torch.zeros_like(accepted_tokens),
    ).contiguous()


def npu_recurrent_gated_delta_rule_310(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor | None,
    beta: torch.Tensor,
    state: torch.Tensor,
    cu_seqlens: torch.Tensor,
    ssm_state_indices: torch.Tensor,
    num_accepted_tokens: torch.Tensor | None = None,
    use_qk_l2norm_in_kernel: bool = True,
) -> torch.Tensor:
    if use_qk_l2norm_in_kernel:
        q = l2norm_310p(q)
        k = l2norm_310p(k)

    total_tokens = v.shape[1]
    flat_state_indices = _flatten_state_indices(ssm_state_indices, cu_seqlens, total_tokens)
    actual_seq_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int32).contiguous()
    flat_state_indices = torch.clamp_min(
        flat_state_indices,
        0,
    ).contiguous()
    accepted_tokens = None
    if num_accepted_tokens is not None:
        accepted_tokens = _mask_padded_recurrent_accepted_tokens(
            num_accepted_tokens,
            actual_seq_lengths,
        )

    out = torch.ops._C_ascend.npu_recurrent_gated_delta_rule_310(
        query=q.squeeze(0).to(torch.float16).contiguous(),
        key=k.squeeze(0).to(torch.float16).contiguous(),
        value=v.squeeze(0).to(torch.float16).contiguous(),
        g=None if g is None else g.squeeze(0).to(torch.float32).contiguous(),
        gk=None,
        beta=beta.squeeze(0).to(torch.float16).contiguous(),
        state=state,
        actual_seq_lengths=actual_seq_lengths,
        ssm_state_indices=flat_state_indices,
        num_accepted_tokens=accepted_tokens,
        scale_value=k.shape[-1] ** -0.5,
    ).unsqueeze(0)
    return out


def _310p_get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
    conv_state_dtype, _ = _original_get_state_dtype(self)
    return conv_state_dtype, torch.float16


_original_get_state_dtype = GatedDeltaNetAttention.get_state_dtype


def _merge_spec_and_non_spec_outputs_310(
    core_attn_out: torch.Tensor,
    num_actual_tokens: int,
    spec_token_indx: torch.Tensor,
    non_spec_token_indx: torch.Tensor,
    core_attn_out_spec: torch.Tensor,
    core_attn_out_non_spec: torch.Tensor,
) -> None:
    """Merge spec/non-spec GDN outputs back into the batch layout.

    Avoid NPU ``index_copy_`` (IndexPutV2) which fails on some layouts; use
    direct indexing instead. Validate lengths so mixed prefill+spec batches
    do not pass mismatched tensors from spec ops.
    """
    spec_out = core_attn_out_spec.squeeze(0)
    non_spec_out = core_attn_out_non_spec.squeeze(0)
    n_spec = spec_token_indx.numel()
    n_non_spec = non_spec_token_indx.numel()
    if spec_out.shape[0] != n_spec:
        raise RuntimeError(f"GDN spec output length {spec_out.shape[0]} != spec_token_indx {n_spec}")
    if non_spec_out.shape[0] != n_non_spec:
        raise RuntimeError(f"GDN non-spec output length {non_spec_out.shape[0]} != non_spec_token_indx {n_non_spec}")
    out = core_attn_out[:num_actual_tokens]
    out[spec_token_indx] = spec_out
    out[non_spec_token_indx] = non_spec_out


class AscendGatedDeltaNetAttention310(GatedDeltaNetAttention):
    get_state_dtype = _310p_get_state_dtype

    def get_attn_backend(self):
        from vllm_ascend._310p.ops.gdn_attn_builder_310 import (
            AscendGDNAttentionBackend310,
        )

        return AscendGDNAttentionBackend310

    def _forward_core(
        self,
        mixed_qkv: torch.Tensor,
        b: torch.Tensor,
        a: torch.Tensor,
        core_attn_out: torch.Tensor,
    ):
        # Core attention computation (called by custom op).

        # NOTE: The processing logic of Qwen3_5GatedDeltaNet is the same as Qwen3NextGatedDeltaNet.
        # However, because the ops `torch_npu.npu_recurrent_gated_delta_rule`
        # currently does not support `ssm_state` inputs in float32 format,
        # we temporarily retain the current _forward_core implementation.
        # Once the ops supports float32 `ssm_state`, this patch should be removed.

        forward_context = get_forward_context()
        attn_metadata: AttentionMetadata = forward_context.attn_metadata

        if attn_metadata is None:
            # V1 profile run
            return

        assert isinstance(attn_metadata, dict)
        attn_metadata = attn_metadata[self.prefix]
        assert isinstance(attn_metadata, GDNAttentionMetadata)
        has_initial_state = attn_metadata.has_initial_state
        spec_query_start_loc = attn_metadata.spec_query_start_loc
        non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
        spec_sequence_masks = attn_metadata.spec_sequence_masks
        spec_token_indx = attn_metadata.spec_token_indx
        non_spec_token_indx = attn_metadata.non_spec_token_indx
        spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor  # noqa: E501
        non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor  # noqa: E501
        self_kv_cache = self.kv_cache
        conv_state = self_kv_cache[0]
        ssm_state = self_kv_cache[1]
        num_actual_tokens = attn_metadata.num_actual_tokens

        if not enable_sp():
            mixed_qkv = mixed_qkv[:num_actual_tokens]
            b = b[:num_actual_tokens]
            a = a[:num_actual_tokens]

        # 1. Convolution sequence transformation
        conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)).transpose(0, 1)
        if spec_sequence_masks is not None:
            if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
                mixed_qkv_spec = mixed_qkv
                mixed_qkv_non_spec = None
            else:
                mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
                mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
        else:
            mixed_qkv_spec = None
            mixed_qkv_non_spec = mixed_qkv
        activation_num = 1 if self.activation else 0

        # 1.1: Process the multi-query part
        if spec_sequence_masks is not None:
            spec_causal_conv1d_meta = attn_metadata.spec_decode_metadata.spec_causal_conv1d
            spec_query_start_loc_device = spec_causal_conv1d_meta.query_start_loc
            uniform_spec_only = attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0
            # The final entry remains the runtime token count even when
            # graph metadata includes padded requests.
            spec_valid_tokens = spec_query_start_loc_device[-1]
            if uniform_spec_only:
                mixed_qkv_spec = _zero_padded_tokens(
                    mixed_qkv_spec,
                    spec_valid_tokens,
                    token_dim=0,
                )
            mixed_qkv_spec = torch.ops._C_ascend.npu_causal_conv1d_310(
                mixed_qkv_spec,
                conv_weights,
                bias=self.conv1d.bias,
                conv_states=conv_state,
                query_start_loc=spec_query_start_loc_device,
                cache_indices=spec_causal_conv1d_meta.cache_indices,
                initial_state_mode=None,
                num_accepted_tokens=spec_causal_conv1d_meta.num_accepted_tokens,
                activation_mode=activation_num,
                pad_slot_id=PAD_SLOT_ID,
                run_mode=1,
            )

        # 1.2: Process the remaining part
        if attn_metadata.num_prefills > 0:
            if mixed_qkv_non_spec is not None:
                mixed_qkv_non_spec = torch.ops._C_ascend.npu_causal_conv1d_310(
                    mixed_qkv_non_spec,
                    conv_weights,
                    bias=self.conv1d.bias,
                    conv_states=conv_state,
                    query_start_loc=non_spec_query_start_loc,
                    cache_indices=non_spec_state_indices_tensor,
                    initial_state_mode=has_initial_state,
                    num_accepted_tokens=None,
                    activation_mode=activation_num,
                    pad_slot_id=PAD_SLOT_ID,
                    run_mode=0,
                )
        elif attn_metadata.num_decodes > 0:
            mixed_qkv_non_spec = torch.ops._C_ascend.npu_causal_conv1d_310(
                mixed_qkv_non_spec,
                conv_weights,
                bias=self.conv1d.bias,
                conv_states=conv_state,
                query_start_loc=None,
                cache_indices=non_spec_state_indices_tensor[: attn_metadata.num_actual_tokens],
                initial_state_mode=None,
                num_accepted_tokens=None,
                activation_mode=activation_num,
                pad_slot_id=PAD_SLOT_ID,
                run_mode=1,
            )
        else:
            mixed_qkv_non_spec = None
        query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
        query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(mixed_qkv_non_spec)

        g, beta = fused_gdn_gating_pytorch(self.A_log, a, b, self.dt_bias)
        if attn_metadata.num_prefills > 0 or spec_sequence_masks is not None:
            if spec_sequence_masks is not None:
                if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
                    g_spec = g
                    beta_spec = beta
                    g_non_spec = None
                    beta_non_spec = None
                else:
                    g_spec = g.index_select(1, spec_token_indx)
                    beta_spec = beta.index_select(1, spec_token_indx)
                    g_non_spec = g.index_select(1, non_spec_token_indx)
                    beta_non_spec = beta.index_select(1, non_spec_token_indx)
            else:
                g_spec = None
                beta_spec = None
                g_non_spec = g
                beta_non_spec = beta

            # 2. Recurrent attention

            # 2.1: Process the multi-query part
            if spec_sequence_masks is not None:
                core_attn_out_spec = npu_recurrent_gated_delta_rule_310(
                    q=query_spec,
                    k=key_spec,
                    v=value_spec,
                    g=g_spec,
                    beta=beta_spec,
                    state=ssm_state,
                    cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],
                    ssm_state_indices=spec_state_indices_tensor,
                    num_accepted_tokens=spec_causal_conv1d_meta.num_accepted_tokens,
                    use_qk_l2norm_in_kernel=True,
                )
            else:
                core_attn_out_spec = None

            # 2.2: Process the remaining part
            if attn_metadata.num_prefills > 0:
                initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
                initial_state[~has_initial_state, ...] = 0
                (
                    core_attn_out_non_spec,
                    last_recurrent_state,
                ) = chunk_gated_delta_rule_310(
                    q=query_non_spec,
                    k=key_non_spec,
                    v=value_non_spec,
                    g=g_non_spec,
                    beta=beta_non_spec,
                    initial_state=initial_state,
                    output_final_state=True,
                    cu_seqlens=non_spec_query_start_loc,
                    head_first=False,
                    use_qk_l2norm_in_kernel=True,
                )

                # Init cache
                ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(ssm_state.dtype)
            elif attn_metadata.num_decodes > 0:
                core_attn_out_non_spec = npu_recurrent_gated_delta_rule_310(
                    q=query_non_spec,
                    k=key_non_spec,
                    v=value_non_spec,
                    g=g_non_spec,
                    beta=beta_non_spec,
                    state=ssm_state,
                    cu_seqlens=non_spec_query_start_loc[: attn_metadata.num_decodes + 1],
                    ssm_state_indices=non_spec_state_indices_tensor,
                    use_qk_l2norm_in_kernel=True,
                )
            else:
                core_attn_out_non_spec = None

        elif attn_metadata.num_decodes > 0:
            core_attn_out_non_spec = npu_recurrent_gated_delta_rule_310(
                q=query_non_spec,
                k=key_non_spec,
                v=value_non_spec,
                g=g,
                beta=beta,
                state=ssm_state,
                cu_seqlens=non_spec_query_start_loc,
                ssm_state_indices=non_spec_state_indices_tensor,
                use_qk_l2norm_in_kernel=True,
            )
        # 3. Merge core attention output
        if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
            _merge_spec_and_non_spec_outputs_310(
                core_attn_out,
                num_actual_tokens,
                spec_token_indx,
                non_spec_token_indx,
                core_attn_out_spec,
                core_attn_out_non_spec,
            )
        elif spec_sequence_masks is not None:
            if not enable_sp():
                core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
            else:
                core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)[:num_actual_tokens]
        else:
            if not enable_sp():
                core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
            else:
                core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)[:num_actual_tokens]
        if spec_sequence_masks is not None and uniform_spec_only:
            core_attn_out.copy_(
                _zero_padded_tokens(
                    core_attn_out,
                    spec_valid_tokens,
                    token_dim=0,
                )
            )
        maybe_save_kv_layer_to_connector("", [])