"""
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# 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 __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, List, Optional

import paddle
from paddle.nn.functional.flash_attention import flash_attn_unpadded
from paddleformers.utils.log import logger

try:
    from paddle.nn.functional.flash_attention import flash_attention_v3_varlen
except Exception as e:
    logger.debug(f"flash_attention_v3_varlen not available: {e}")
    flash_attention_v3_varlen = None

try:
    from paddle.nn.functional.flash_attention import flashmask_attention
except Exception as e:
    logger.debug(f"flashmask_attention not available: {e}")
    flashmask_attention = None

from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
    AttentionBackend,
    AttentionMetadata,
)
from fastdeploy.model_executor.layers.attention.ops import (
    append_attention,
    get_attn_mask_q,
    get_block_shape_and_split_kv_block,
    gqa_rope_write_cache,
    init_kv_signal_per_query,
    init_signal_layerwise,
    open_shm_and_get_meta_signal,
    pre_cache_len_concat,
)
from fastdeploy.model_executor.layers.attention.utils import init_rank_and_device_id
from fastdeploy.model_executor.utils import get_sm_version

if TYPE_CHECKING:
    from fastdeploy.model_executor.forward_meta import ForwardMeta

import os

from fastdeploy import envs
from fastdeploy.platforms import current_platform

flashmask_attention_v4 = None

if current_platform.is_cuda():
    from fastdeploy.model_executor.ops.gpu import merge_prefill_decode_output
else:
    merge_prefill_decode_output = None


from fastdeploy.spec_decode import SpecMethod

FLASH_ATTN_VERSION = None


def init_flash_attn_version():
    """
    init_flash_attn_version
    """
    if current_platform.is_cuda():
        global FLASH_ATTN_VERSION
        sm_version = get_sm_version()
        if sm_version >= 100:
            try:
                paddle.enable_compat(scope={"cutlass"})
                from flash_mask.cute.interface import flashmask_attention as fa4

                global flashmask_attention_v4
                flashmask_attention_v4 = fa4
                FLASH_ATTN_VERSION = 4
                logger.info("The current platform supports Flash Attention V4.")
            except ImportError:
                logger.info(f"The current platform[sm{get_sm_version()}] can't import Flash Attention V4.")

        if FLASH_ATTN_VERSION is None:
            if sm_version == 90 and 90 in paddle.version.cuda_archs():
                FLASH_ATTN_VERSION = 3
                logger.info("The current platform supports Flash Attention V3.")
            else:
                FLASH_ATTN_VERSION = 2
                logger.info("The current platform only support Flash Attention V2.")
    else:
        logger.info("Only support CUDA version flash attention.")


def _is_deterministic_mode():
    """Check if FD_DETERMINISTIC_MODE is enabled."""
    return envs.FD_DETERMINISTIC_MODE


init_flash_attn_version()


def flash_attn_func(
    q: paddle.Tensor,
    k: paddle.Tensor,
    v: paddle.Tensor,
    cu_seqlens_q: Optional[paddle.Tensor] = None,
    cu_seqlens_k: Optional[paddle.Tensor] = None,
    max_seqlen_q: Optional[paddle.Tensor] = None,
    max_seqlen_k: Optional[paddle.Tensor] = None,
    attn_mask_q: Optional[paddle.Tensor] = None,
    causal: bool = True,
    num_heads: int = None,
    kv_num_heads: int = None,
    head_dim: int = 128,
    version: Optional[int] = None,
):
    if FLASH_ATTN_VERSION is None:
        init_flash_attn_version()
    if version is None:
        version = FLASH_ATTN_VERSION

    if version == 4:
        assert (
            flashmask_attention_v4 is not None
        ), "Cannot import flashmask_attention from flash_mask.cute.interface, please install it first"
        assert attn_mask_q is not None, "FA4 requires attn_mask_q"
        assert num_heads is not None
        assert kv_num_heads is not None
        original_flash_attn_version = paddle.base.framework.get_flags(["FLAGS_flash_attn_version"])[
            "FLAGS_flash_attn_version"
        ]
        with paddle.no_grad():
            try:
                paddle.set_flags({"FLAGS_flash_attn_version": 4})
                out = flashmask_attention_v4(
                    q.reshape([1, -1, num_heads, head_dim]),
                    k.reshape([1, -1, kv_num_heads, head_dim]),
                    v.reshape([1, -1, kv_num_heads, head_dim]),
                    startend_row_indices=attn_mask_q,
                    causal=False,
                    return_softmax_lse=True,
                    training=True,
                )
            finally:
                paddle.set_flags({"FLAGS_flash_attn_version": original_flash_attn_version})
        return out

    if attn_mask_q is not None:
        assert flashmask_attention is not None
        out = flashmask_attention(
            q.reshape([1, -1, num_heads, head_dim]),
            k.reshape([1, -1, kv_num_heads, head_dim]),
            v.reshape([1, -1, kv_num_heads, head_dim]),
            startend_row_indices=attn_mask_q,
            causal=False,
        )
    else:
        if version == 3:
            out = flash_attention_v3_varlen(
                q,
                k,
                v,
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_k,
                max_seqlen_q=max_seqlen_q,
                max_seqlen_k=max_seqlen_k,
                causal=causal,
            )
        else:
            out = flash_attn_unpadded(
                q,
                k,
                v,
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_k,
                max_seqlen_q=max_seqlen_q,
                max_seqlen_k=max_seqlen_k,
                causal=causal,
                scale=head_dim**-0.5,
                training=False,
            )
    return out


@dataclass
class FlashAttentionMetadata(AttentionMetadata):
    """
    FlashAttentionMetadata
    """

    cu_seqlens_k: paddle.Tensor = None

    pre_cache_batch_ids = None
    pre_cache_tile_ids_per_batch = None
    pre_cache_num_blocks_cpu = None
    kv_token_num_cpu = None

    # pd_disaggregation
    kv_signal_metadata: Optional[paddle.Tensor] = None
    kv_signal_data_list: List[Optional[paddle.Tensor]] = field(default_factory=list)

    _fuse_kernel_compute_dtype: str = "bf16"
    _dtype: paddle.dtype = paddle.bfloat16

    max_len_tensor_cpu_decoder: paddle.Tensor = None

    attn_mask_q: paddle.Tensor = None


class FlashAttentionBackend(AttentionBackend):
    """
    FlashAttentionBackend backend implementation
    """

    __infer_dynamic_dims_fields__ = ["attention_metadata"]
    attention_metadata: FlashAttentionMetadata

    def __init__(
        self,
        fd_config: FDConfig,
        kv_num_heads: int,
        num_heads: int,
        head_dim: int,
        encoder_block_shape_q: int = -1,
        decoder_block_shape_q: int = -1,
    ):
        """
        FlashAttentionBackend __init__
        """
        super().__init__()
        self.max_seq_len = fd_config.model_config.max_model_len
        self.causal = getattr(fd_config.model_config, "causal", True)

        self.kv_num_heads = kv_num_heads
        self.num_heads = num_heads
        self.group_size: int = self.num_heads // self.kv_num_heads
        self.head_dim = fd_config.model_config.head_dim
        self.attn_outputsize_tp = self.num_heads * self.head_dim
        self.block_size = fd_config.cache_config.block_size
        self.num_layers: int = fd_config.model_config.num_hidden_layers
        self.encoder_block_shape_q: int = encoder_block_shape_q
        self.decoder_block_shape_q: int = decoder_block_shape_q

        self.speculative_method = fd_config.speculative_config.method
        self.use_speculate = self.speculative_method is not None
        self.speculate_max_draft_token_num = fd_config.speculative_config.num_speculative_tokens
        if not self.use_speculate:
            self.speculate_max_draft_token_num = 0
        self.keep_pd_step_flag: bool = fd_config.speculative_config.model_type == "mtp"
        self.num_layers_draft_model: int = int(fd_config.speculative_config.method == SpecMethod.MTP)

        self.pd_disaggregation_mode: str = fd_config.parallel_config.pd_disaggregation_mode

        self.start_layer_index: int = fd_config.model_config.start_layer_index

        self.rank, self.device_id = init_rank_and_device_id(fd_config)

        self.rope_3d: bool = fd_config.enable_rope_3d_runtime
        if fd_config.speculative_config.model_type != "main":
            self.rope_3d = False
        # Note(ZKK): here must be consistent with append_attn_backend.py
        self.max_partition_size: int = int(os.getenv("FLAGS_max_partition_size", 1024))
        if FLASH_ATTN_VERSION is None:
            init_flash_attn_version()

    def get_attention_meta(self):
        """get_attention_meta"""
        return self.attention_metadata

    def get_kv_cache_shape(
        self,
        max_num_blocks: int,
        kv_cache_quant_type: str = None,
    ):
        """
        Calculate kv cache shape
        """
        key_cache_shape = [max_num_blocks, self.kv_num_heads, self.block_size, self.head_dim]
        if kv_cache_quant_type is not None and kv_cache_quant_type == "int4_zp":
            key_cache_shape[-1] = self.head_dim // 2
        value_cache_shape = key_cache_shape
        return key_cache_shape, value_cache_shape

    def init_attention_metadata(self, forward_meta: ForwardMeta):
        metadata = FlashAttentionMetadata()

        # pd_disaggregation
        metadata.kv_signal_data_list = [None] * self.num_layers
        if self.pd_disaggregation_mode == "per_chunk":
            if not self.keep_pd_step_flag and not forward_meta.is_dummy_or_profile_run:
                init_kv_signal_per_query(
                    forward_meta.seq_lens_encoder,
                    forward_meta.seq_lens_this_time,
                    forward_meta.seq_lens_decoder,
                    self.rank,
                    self.num_layers + self.num_layers_draft_model,
                )
        elif self.pd_disaggregation_mode == "per_query":
            metadata.kv_signal_metadata = open_shm_and_get_meta_signal(
                self.rank, int(self.device_id), self.keep_pd_step_flag
            )

        if metadata._dtype == "bfloat16":
            metadata._fuse_kernel_compute_dtype = "bf16"
        elif metadata._dtype == "float16":
            metadata._fuse_kernel_compute_dtype = "fp16"
        elif metadata._dtype == "float32":
            metadata._fuse_kernel_compute_dtype = "fp32"

        self.attention_metadata = metadata

    def forward_mixed(
        self,
        q: paddle.Tensor,
        k: paddle.Tensor,
        v: paddle.Tensor,
        qkv: paddle.Tensor,
        compressed_kv: paddle.Tensor,
        k_pe: paddle.Tensor,
        layer: Attention,
        forward_meta: ForwardMeta,
    ):
        metadata = self.attention_metadata

        if self.pd_disaggregation_mode == "per_query":
            metadata.kv_signal_data_list[layer.layer_id] = init_signal_layerwise(
                metadata.kv_signal_metadata,
                layer.layer_id + self.start_layer_index,
            )

        if int(os.getenv("USE_TBO", "0")) == 1:
            if hasattr(forward_meta, "tbo_microbatch_id"):
                # here we only let the last microbatch invoke cache kv transfer！
                if forward_meta.tbo_microbatch_id == 0:
                    os.environ["FLAGS_fmt_write_cache_completed_signal"] = "0"
                elif forward_meta.tbo_microbatch_id == 1:
                    os.environ["FLAGS_fmt_write_cache_completed_signal"] = "1"

        norm_after_rope_in_kernel = not getattr(layer, "qk_norm_before_rope", False)
        q_norm_weight = getattr(layer, "q_norm_weight", None) if norm_after_rope_in_kernel else None
        k_norm_weight = getattr(layer, "k_norm_weight", None) if norm_after_rope_in_kernel else None

        cache_quant_type_str = getattr(layer, "cache_quant_type_str", "none")
        if cache_quant_type_str == "block_wise_fp8":
            cache_k = forward_meta.caches[4 * layer.layer_id]
            cache_v = forward_meta.caches[4 * layer.layer_id + 1]
            cache_k_scales = forward_meta.caches[4 * layer.layer_id + 2]
            cache_v_scales = forward_meta.caches[4 * layer.layer_id + 3]
        else:
            cache_k = forward_meta.caches[2 * layer.layer_id]
            cache_v = forward_meta.caches[2 * layer.layer_id + 1]
            cache_k_scales = getattr(layer, "cache_k_scale", None)
            cache_v_scales = getattr(layer, "cache_v_scale", None)

        if layer.layer_id == 0:
            get_block_shape_and_split_kv_block(
                forward_meta.seq_lens_encoder,
                forward_meta.seq_lens_decoder,
                forward_meta.seq_lens_this_time,
                forward_meta.decoder_batch_ids,
                forward_meta.decoder_tile_ids_per_batch,
                forward_meta.decoder_num_blocks_cpu,
                forward_meta.decoder_num_blocks_device,
                forward_meta.decoder_chunk_size_device,
                forward_meta.max_len_tensor_cpu,
                forward_meta.encoder_batch_ids,
                forward_meta.encoder_tile_ids_per_batch,
                forward_meta.encoder_num_blocks_x_cpu,
                forward_meta.kv_batch_ids,
                forward_meta.kv_tile_ids_per_batch,
                forward_meta.kv_num_blocks_x_cpu,
                self.encoder_block_shape_q,
                self.decoder_block_shape_q,
                self.group_size,
                self.block_size,
            )

            if forward_meta.max_len_tensor_cpu[1].item() > 0:

                forward_meta.max_len_tensor_cpu_decoder = paddle.clone(forward_meta.max_len_tensor_cpu)
                forward_meta.max_len_tensor_cpu_decoder[1] = 0

                (
                    forward_meta.cu_seqlens_k,
                    forward_meta.pre_cache_batch_ids,
                    forward_meta.pre_cache_tile_ids_per_batch,
                    forward_meta.pre_cache_num_blocks_cpu,
                    forward_meta.kv_token_num_cpu,
                ) = pre_cache_len_concat(
                    forward_meta.seq_lens_encoder,
                    forward_meta.seq_lens_decoder,
                    forward_meta.seq_lens_this_time,
                    forward_meta.max_len_tensor_cpu[2],
                    self.block_size,
                )
                if FLASH_ATTN_VERSION == 4 or forward_meta.attn_mask_offsets is not None:
                    forward_meta.attn_mask_q = get_attn_mask_q(
                        cu_seqlens_q=forward_meta.cu_seqlens_q,
                        cu_seqlens_k=forward_meta.cu_seqlens_k,
                        attn_mask_kv=forward_meta.attn_mask_offsets,
                        kv_token_num=forward_meta.kv_token_num_cpu[0].item(),
                    )
                else:
                    forward_meta.attn_mask_q = None

        use_fa_do_prefill = forward_meta.max_len_tensor_cpu[1].item() > 0

        if use_fa_do_prefill:
            q, k, v, _ = gqa_rope_write_cache(
                qkv,
                cache_k,
                cache_v,
                forward_meta.cu_seqlens_q,
                forward_meta.cu_seqlens_k,
                forward_meta.rotary_embs,
                forward_meta.seq_lens_this_time,
                forward_meta.seq_lens_encoder,
                forward_meta.seq_lens_decoder,
                forward_meta.batch_id_per_token,
                forward_meta.block_tables,
                forward_meta.kv_batch_ids,
                forward_meta.kv_tile_ids_per_batch,
                forward_meta.kv_num_blocks_x_cpu,
                forward_meta.pre_cache_batch_ids,
                forward_meta.pre_cache_tile_ids_per_batch,
                forward_meta.pre_cache_num_blocks_cpu,
                q_norm_weight,
                k_norm_weight,
                cache_k_scales,
                cache_v_scales,
                getattr(layer, "cache_k_out_scale", None),
                getattr(layer, "cache_v_out_scale", None),
                getattr(layer, "cache_k_zp", None),
                getattr(layer, "cache_v_zp", None),
                metadata.kv_signal_data_list[layer.layer_id],
                forward_meta.kv_token_num_cpu[0].item(),
                self.max_seq_len,
                getattr(layer, "rms_norm_eps", 1e-6),
                layer.use_neox_rotary_style,
                getattr(layer, "cache_quant_type_str", "none"),
                self.rope_3d,
            )

            res_encoder = flash_attn_func(
                q,
                k,
                v,
                forward_meta.cu_seqlens_q[: forward_meta.cu_seqlens_k.shape[0]],
                forward_meta.cu_seqlens_k,
                max_seqlen_q=forward_meta.max_len_tensor_cpu[0],
                max_seqlen_k=forward_meta.max_len_tensor_cpu[3],
                attn_mask_q=forward_meta.attn_mask_q,
                causal=self.causal,
                num_heads=self.num_heads,
                kv_num_heads=self.kv_num_heads,
                head_dim=self.head_dim,
            )[0].reshape([-1, self.attn_outputsize_tp])

        res_decoder = append_attention(
            qkv,
            cache_k,
            cache_v,
            forward_meta.seq_lens_encoder,
            forward_meta.seq_lens_decoder,
            forward_meta.seq_lens_this_time,
            forward_meta.batch_id_per_token,
            forward_meta.cu_seqlens_q,
            forward_meta.block_tables,
            forward_meta.encoder_batch_ids,
            forward_meta.encoder_tile_ids_per_batch,
            forward_meta.encoder_num_blocks_x_cpu,
            forward_meta.kv_batch_ids,
            forward_meta.kv_tile_ids_per_batch,
            forward_meta.kv_num_blocks_x_cpu,
            forward_meta.decoder_batch_ids,
            forward_meta.decoder_tile_ids_per_batch,
            forward_meta.decoder_num_blocks_cpu,
            forward_meta.max_len_tensor_cpu_decoder if use_fa_do_prefill else forward_meta.max_len_tensor_cpu,
            forward_meta.rotary_embs,
            forward_meta.attn_mask,
            layer.qkv_bias,
            layer.qkv_scale,
            cache_k_scales,
            cache_v_scales,
            getattr(layer, "cache_k_out_scale", None),
            getattr(layer, "cache_v_out_scale", None),
            getattr(layer, "cache_k_zp", None),
            getattr(layer, "cache_v_zp", None),
            layer.linear_shift,
            layer.linear_smooth,
            forward_meta.attn_mask_offsets,
            metadata.kv_signal_data_list[layer.layer_id],
            q_norm_weight,
            k_norm_weight,
            getattr(layer, "sinks", None),
            getattr(layer, "rms_norm_eps", 1e-6),
            metadata._fuse_kernel_compute_dtype,
            getattr(layer, "cache_quant_type_str", "none"),
            layer.use_neox_rotary_style,
            self.rope_3d,
            self.max_seq_len,
            getattr(layer, "quant_max_bound", 0.0),
            getattr(layer, "quant_min_bound", 0.0),
            getattr(layer, "out_scale", -1.0),
            self.encoder_block_shape_q,
            self.decoder_block_shape_q,
            self.max_partition_size,
            self.max_seq_len,
            self.speculate_max_draft_token_num + 1,
            self.causal,
            self.speculative_method is not None,
        )

        if use_fa_do_prefill:
            merge_prefill_decode_output(
                res_encoder,
                res_decoder,
                forward_meta.seq_lens_encoder,
                forward_meta.seq_lens_decoder,
                forward_meta.seq_lens_this_time,
                forward_meta.cu_seqlens_q,
                self.num_heads,
                self.head_dim,
                self.speculate_max_draft_token_num + 1,
            )
            return res_encoder
        else:
            return res_decoder
