"""
Copyright (c) 2025 by FlashInfer team.

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.
"""

import functools
import math
from typing import Optional, Tuple, Union

import torch

from ..api_logging import flashinfer_api
from ..jit import gen_batch_attention_module
from ..trace.templates.attention import batch_attention_run_trace
from ..utils import (
    MaskMode,
    PosEncodingMode,
    TensorLayout,
    _check_kv_layout,
    _unpack_paged_kv_cache,
    determine_attention_backend,
)
from ..prefill import BatchPrefillWithPagedKVCacheWrapper
from ..jit.attention.variants import attention_sink_decl
from ..jit.utils import filename_safe_dtype_map


@functools.cache
def get_holistic_attention_module(*args):
    return gen_batch_attention_module(*args).build_and_load()


class BatchAttention:
    r"""Holistic batched attention wrapper that fuses paged-prefill and paged-decode requests
    into a single kernel launch.

    ``BatchAttention`` dispatches between prefill-style and decode-style execution per
    request based on the ``qo_indptr`` / ``kv_indptr`` ranges supplied to :meth:`plan`, so a
    serving stack can submit a mixed batch (e.g. some prompts in prefill, others in decode)
    without splitting it into two separate wrappers.  Workspace buffers are owned by the
    instance and reused across :meth:`plan` / :meth:`run` calls.

    Parameters
    ----------
    kv_layout : str
        Layout of the paged KV-cache tensors, either ``"NHD"`` (token-major) or
        ``"HND"`` (head-major).  Defaults to ``"NHD"``.
    device : str
        CUDA device that owns the internal workspace buffers, e.g. ``"cuda"`` or
        ``"cuda:0"``.  Defaults to ``"cuda"``.
    """

    @flashinfer_api
    def __init__(
        self,
        kv_layout: str = "NHD",
        device: str = "cuda",
    ):
        r"""Allocate workspace buffers and bind the wrapper to a CUDA device.

        See :class:`BatchAttention` for the meaning of each parameter.
        """
        _check_kv_layout(kv_layout)
        self._kv_layout = kv_layout

        self.float_workspace_buffer = torch.empty(
            384 * 1024 * 1024,
            dtype=torch.uint8,
            device=torch.device(device),
        )
        self.int_workspace_buffer = torch.empty(
            8 * 1024 * 1024,
            dtype=torch.uint8,
            device=torch.device(device),
        )
        self.page_locked_int_workspace_buffer = torch.empty(
            8 * 1024 * 1024,
            dtype=torch.uint8,
            device=torch.device("cpu"),
            pin_memory=True,
        )

    @flashinfer_api
    def plan(
        self,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        kv_indices: torch.Tensor,
        kv_len_arr: torch.Tensor,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim_qk: int,
        head_dim_vo: int,
        page_size: int,
        causal: bool = False,
        sm_scale: float = None,
        logits_soft_cap: Optional[float] = None,
        q_data_type: torch.dtype = torch.bfloat16,
        kv_data_type: torch.dtype = torch.bfloat16,
        use_profiler: bool = False,
    ) -> None:
        r"""Plan the holistic attention kernel for a specific batch shape.

        Should be called before any :meth:`run` call.  The plan is cached on the
        instance and reused across subsequent :meth:`run` invocations with the same
        layout.

        Parameters
        ----------
        qo_indptr : torch.Tensor
            CSR-style query offsets, shape ``[batch_size + 1]``, dtype ``int32``.
        kv_indptr : torch.Tensor
            CSR-style page offsets into ``kv_indices``, shape ``[batch_size + 1]``,
            dtype ``int32``.
        kv_indices : torch.Tensor
            Page indices into the paged KV-cache, shape ``[kv_indptr[-1]]``,
            dtype ``int32``.
        kv_len_arr : torch.Tensor
            Per-request KV-cache lengths in tokens, shape ``[batch_size]``,
            dtype ``int32``.
        num_qo_heads : int
            Number of query / output heads.
        num_kv_heads : int
            Number of key / value heads.  Must divide ``num_qo_heads``.
        head_dim_qk : int
            Per-head dimension of the query / key tensors.
        head_dim_vo : int
            Per-head dimension of the value / output tensors.
        page_size : int
            Page size of the paged KV-cache.
        causal : bool
            Whether to apply a causal mask.  Defaults to ``False``.
        sm_scale : float
            Softmax scale.  If ``None``, defaults to ``1/sqrt(head_dim_qk)``.
        logits_soft_cap : Optional[float]
            Logits soft-cap value.  ``None`` or ``0`` disables capping.
        q_data_type : torch.dtype
            Dtype of the query tensor.  Defaults to ``torch.bfloat16``.
        kv_data_type : torch.dtype
            Dtype of the key / value tensors.  Defaults to ``torch.bfloat16``.
        use_profiler : bool
            Whether to compile the profiler-enabled variant of the kernel.  Defaults
            to ``False``.
        """
        if logits_soft_cap is None:
            logits_soft_cap = 0.0
        self._logits_soft_cap = logits_soft_cap

        # head_dim > 256 is for holistic (persistent) kernel.
        if head_dim_qk > 256 or head_dim_vo > 256:
            raise ValueError(
                "BatchAttention (holistic persistent kernel) does not support "
                f"head_dim > 256 (got head_dim_qk={head_dim_qk}, "
                f"head_dim_vo={head_dim_vo}). Use "
                "BatchPrefillWithPagedKVCacheWrapper(backend='fa2') or "
                "BatchDecodeWithPagedKVCacheWrapper(use_tensor_cores=True) instead."
            )

        # get jit module
        get_module_args = (
            q_data_type,
            kv_data_type,
            q_data_type,
            kv_indptr.dtype,
            head_dim_qk,
            head_dim_vo,
            PosEncodingMode["NONE"].value,
            logits_soft_cap > 0.0,
            use_profiler,  # different compiler path
        )
        self.module = get_holistic_attention_module(*get_module_args)

        qo_indptr_host = qo_indptr.to(torch.device("cpu"), non_blocking=True)
        kv_indptr_host = kv_indptr.to(torch.device("cpu"), non_blocking=True)
        kv_len_arr_host = kv_len_arr.to(torch.device("cpu"), non_blocking=True)
        torch.cuda.synchronize()

        batch_size = kv_len_arr.shape[0]
        self._page_size = page_size
        self._sm_scale = sm_scale
        self._mask_mode = MaskMode.CAUSAL.value if causal else MaskMode.NON_CAUSAL.value
        self._num_qo_heads = num_qo_heads
        self._num_kv_heads = num_kv_heads
        self._page_size = page_size
        self._use_profiler = use_profiler

        # No addtional buf allocated for CUDA graph tensor
        # Allocate outside FlashInfer
        self._kv_indices = kv_indices
        self._plan_info = self.module.plan(
            self.float_workspace_buffer,
            self.int_workspace_buffer,
            self.page_locked_int_workspace_buffer,
            qo_indptr_host,
            kv_indptr_host,
            kv_len_arr_host,
            batch_size,
            num_qo_heads,
            num_kv_heads,
            head_dim_vo,
            causal,
        )

    @flashinfer_api(trace=batch_attention_run_trace)
    def run(
        self,
        q: torch.Tensor,
        kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        k_scale: Optional[torch.Tensor] = None,
        v_scale: Optional[torch.Tensor] = None,
        logits_soft_cap: float = 0.0,
        profiler_buffer: Optional[torch.Tensor] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""Run the planned holistic attention kernel.

        Parameters
        ----------
        q : torch.Tensor
            Query tensor, shape ``[total_qo_tokens, num_qo_heads, head_dim_qk]``.
        kv_cache : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            Either a single packed paged KV-cache tensor (when K and V share storage) or a
            ``(k_cache, v_cache)`` pair.  Layout must match the ``kv_layout`` passed to
            :meth:`__init__`.
        out : Optional[torch.Tensor]
            Optional output buffer.  If ``None``, a new tensor is allocated with the same
            shape as ``q``.
        lse : Optional[torch.Tensor]
            Optional log-sum-exp buffer, shape ``[total_qo_tokens, num_qo_heads]``, dtype
            ``float32``.  Allocated if ``None``.
        k_scale : Optional[torch.Tensor]
            FP8 dequantization scale for ``k``.  Pre-multiplied into ``sm_scale``.
        v_scale : Optional[torch.Tensor]
            FP8 dequantization scale for ``v``.  Applied to the output.
        logits_soft_cap : float
            Logits soft-cap value.  Must be consistent with the ``logits_soft_cap``
            passed to :meth:`plan` (a non-zero value here requires a non-zero plan-time
            value too).
        profiler_buffer : Optional[torch.Tensor]
            Profiler buffer.  Required if the wrapper was planned with
            ``use_profiler=True``.
        kv_cache_sf : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]
            Optional scale tensors for NVFP4 KV-cache (one tensor or a ``(k_sf, v_sf)``
            pair, mirroring the structure of ``kv_cache``).

        Returns
        -------
        Tuple[torch.Tensor, torch.Tensor]
            ``(out, lse)`` — the attention output and its log-sum-exp.
        """
        if profiler_buffer is None:
            if self._use_profiler:
                raise ValueError(
                    "Profiler is enabled, profiler_buffer must be provided"
                )
        if logits_soft_cap > 0.0 and self._logits_soft_cap <= 0.0:
            raise ValueError(
                "logits_soft_cap used in kernel run but not provided in plan(). This will cause template deduction error."
            )

        k_cache, v_cache = _unpack_paged_kv_cache(kv_cache, self._kv_layout)
        if out is None:
            out = torch.empty_like(q)
        if lse is None:
            # lse shape: [batch_size, num_qo_heads]
            lse = torch.empty(
                q.shape[0], q.shape[1], device=q.device, dtype=torch.float32
            )
        head_dim_qk = q.shape[2]
        sm_scale = self._sm_scale
        if sm_scale is None:
            sm_scale = 1.0 / math.sqrt(head_dim_qk)
        if k_scale is not None:
            sm_scale *= k_scale
        if v_scale is None:
            v_scale = 1.0
        # profiler_buffer is optional
        profiler_args = (profiler_buffer,) if self._use_profiler else ()

        # Unpack kv_cache_sf for NVFP4 (maybe_k_cache_sf, maybe_v_cache_sf)
        k_cache_sf, v_cache_sf = (
            _unpack_paged_kv_cache(kv_cache_sf, self._kv_layout)
            if kv_cache_sf is not None
            else (None, None)
        )

        self.module.run(
            self.float_workspace_buffer,
            self.int_workspace_buffer,
            self._plan_info,
            q,
            k_cache,
            v_cache,
            self._kv_indices,
            out,
            lse,
            self._mask_mode,
            TensorLayout[self._kv_layout].value,
            self._num_qo_heads,
            self._num_kv_heads,
            self._page_size,
            v_scale,
            sm_scale,
            logits_soft_cap,
            # ADDITIONAL_FUNC_PARAMS (maybe_k_cache_sf, maybe_v_cache_sf)
            k_cache_sf,
            v_cache_sf,
            # PROFILER_FUNC_PARAMS
            *profiler_args,
        )

        return out, lse


class BatchAttentionWithAttentionSinkWrapper(BatchPrefillWithPagedKVCacheWrapper):
    r"""
    Wrapper for prefill and decode attention with paged KV-cache that adds support for
    attention sinks. This class extends `BatchPrefillWithPagedKVCacheWrapper`, providing
    a convenient interface for using attention sinks during prefill or decode attention.
    """

    # No @flashinfer_api here: parent class BatchPrefillWithPagedKVCacheWrapper
    # already decorates __init__, so decorating again produces double log entries.
    def __init__(
        self,
        float_workspace_buffer: torch.Tensor,
        kv_layout: str = "NHD",
        use_cuda_graph: bool = False,
        qo_indptr_buf: Optional[torch.Tensor] = None,
        paged_kv_indptr_buf: Optional[torch.Tensor] = None,
        paged_kv_indices_buf: Optional[torch.Tensor] = None,
        paged_kv_last_page_len_buf: Optional[torch.Tensor] = None,
        custom_mask_buf: Optional[torch.Tensor] = None,
        mask_indptr_buf: Optional[torch.Tensor] = None,
        backend: str = "auto",
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        q_data_type: torch.dtype = torch.bfloat16,
        kv_data_type: torch.dtype = torch.bfloat16,
        head_dim_qk: int = 128,
        head_dim_vo: int = 128,
        window_left: int = -1,
    ) -> None:
        # trtllm is separate code path
        assert backend in ["fa2", "fa3", "auto"]
        if backend == "auto":
            # dispatch backend before init jit module
            backend = determine_attention_backend(
                float_workspace_buffer.device,
                PosEncodingMode[pos_encoding_mode].value,
                use_fp16_qk_reduction,  # use_fp16_qk_reduction
                custom_mask_buf is not None,  # use_custom_mask
                q_data_type,
                kv_data_type,
                head_dim_qk=head_dim_qk,
                head_dim_vo=head_dim_vo,
            )

        jit_args = [
            f"batch_prefill_attention_sink_{filename_safe_dtype_map[q_data_type]}_swa_{window_left >= 0}_{backend}",  # uri
            q_data_type,  # dtype_q
            kv_data_type,  # dtype_kv
            q_data_type,  # dtype_o
            torch.int32,  # idtype
            head_dim_qk,  # hidden_dim_qk
            head_dim_vo,  # hidden_dim_vo
            ["sink"],  # additional_tensor_names
            ["float"],  # additional_tensor_dtypes
            ["sm_scale"],  # additional_scalar_names
            ["double"],  # additional_scalar_dtypes
            "AttentionSink",
            attention_sink_decl[backend],
        ]
        jit_kwargs = {
            "use_sliding_window": window_left >= 0,
            "use_fp16_qk_reduction": use_fp16_qk_reduction,
            "pos_encoding_mode": PosEncodingMode[pos_encoding_mode].value,
        }

        super().__init__(
            float_workspace_buffer=float_workspace_buffer,
            kv_layout=kv_layout,
            use_cuda_graph=use_cuda_graph,
            qo_indptr_buf=qo_indptr_buf,
            paged_kv_indptr_buf=paged_kv_indptr_buf,
            paged_kv_indices_buf=paged_kv_indices_buf,
            paged_kv_last_page_len_buf=paged_kv_last_page_len_buf,
            custom_mask_buf=custom_mask_buf,
            mask_indptr_buf=mask_indptr_buf,
            backend=backend,
            jit_args=jit_args,
            jit_kwargs=jit_kwargs,
        )
