# Copyright (c) 2026 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.
"""
CuTe DSL FMHA Kernel Loader (Cubin Distribution)
=================================================

Loads pre-compiled FMHA kernel .so artifacts via ExternalBinaryModule.
The .so files are compiled offline from the proprietary DSL kernel source
and distributed through the cubin_publishing pipeline.

Runtime flow:
    1. get_artifact() downloads .so from artifactory (or local cache)
    2. ExternalBinaryModule loads .so and extracts callable kernel
    3. cute_dsl_fmha_ragged_prefill() wraps the kernel with a PyTorch-friendly API
"""

import functools
import logging
import math
import os
from typing import Optional

import cutlass
import cutlass.cute as cute
import torch
from cuda.bindings import driver as cuda_driver
from cutlass.cute.runtime import from_dlpack
from cutlass.cute.typing import Float32, Int32

logger = logging.getLogger("flashinfer.attention.cute_dsl.fmha")


# =============================================================================
# Artifact configuration
# =============================================================================

from flashinfer.artifacts import ArtifactPath

DSL_FMHA_ARTIFACT_PATH = ArtifactPath.DSL_FMHA

from flashinfer.artifacts import _get_host_cpu_arch as _get_cpu_arch


def _get_gpu_arch(device: torch.device) -> str:
    """Get the GPU architecture string for *device* (e.g. 'sm_100a').

    Uses the same normalization as CompilationContext to match the arch
    subdirectory names used by the cubin publishing pipeline.
    """
    major, minor = torch.cuda.get_device_capability(device)
    from flashinfer.compilation_context import CompilationContext

    _, minor_str = CompilationContext._normalize_cuda_arch(major, minor)
    return f"sm_{major}{minor_str}"


# Cache: arch -> {variant_name: sha256} parsed from downloaded checksums.txt
_checksums_cache: dict[str, dict[str, str]] = {}


def _get_checksums(gpu_arch: str) -> dict[str, str]:
    """Download and parse checksums.txt for the given arch.

    The checksums.txt is generated by compile_dsl_fmha.py and published
    alongside the .so files at <artifact_path>/<arch>/checksums.txt.
    Format per line: "sha256  filename.so"
    """
    if gpu_arch in _checksums_cache:
        return _checksums_cache[gpu_arch]

    from flashinfer.jit.cubin_loader import download_file, safe_urljoin
    from flashinfer.jit.cubin_loader import FLASHINFER_CUBINS_REPOSITORY
    from flashinfer.jit.env import FLASHINFER_CUBIN_DIR

    cpu_arch = _get_cpu_arch()
    checksums_rel = os.path.join(
        DSL_FMHA_ARTIFACT_PATH, cpu_arch, gpu_arch, "checksums.txt"
    )
    local_path = FLASHINFER_CUBIN_DIR / checksums_rel

    if not local_path.exists():
        os.makedirs(local_path.parent, exist_ok=True)
        uri = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, checksums_rel)
        logger.info(f"Fetching checksums for {gpu_arch} from {uri}")
        download_file(uri, str(local_path))

    checksums = {}
    with open(local_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            sha256, filename = line.split(None, 1)
            variant_name = filename.rsplit(".", 1)[0]  # strip .so extension
            checksums[variant_name] = sha256

    _checksums_cache[gpu_arch] = checksums
    return checksums


def _dtype_to_str(dtype: torch.dtype) -> str:
    """Convert torch dtype to short string matching compile_cute_dsl_fmha.py naming."""
    return {
        torch.float16: "fp16",
        torch.bfloat16: "bf16",
        torch.float8_e4m3fn: "e4m3",
    }[dtype]


def _get_variant_name(
    qk_dtype: torch.dtype,
    pv_dtype: torch.dtype,
    out_dtype: torch.dtype,
    head_dim: int,
    is_causal: bool,
    is_persistent: bool = True,
    varlen: bool = False,
    with_lse: bool = False,
    enable_tvm_ffi: bool = False,
    enable_skip_softmax: bool = False,
    enable_sink: bool = False,
    use_pdl: bool = False,
) -> str:
    """Generate the variant name matching compile_cute_dsl_fmha.py naming convention."""
    qk_str = _dtype_to_str(qk_dtype)
    pv_str = _dtype_to_str(pv_dtype)
    out_str = _dtype_to_str(out_dtype)
    # Include all three dtypes for partially quantized kernels.  Preserve the
    # compact legacy names when QK and PV use the same dtype.
    dtype_str = (
        f"{qk_str}_{pv_str}_{out_str}"
        if qk_dtype != pv_dtype
        else f"{qk_str}_{out_str}"
        if qk_dtype != out_dtype
        else qk_str
    )
    causal_str = "causal" if is_causal else "nocausal"
    persist_str = "persistent" if is_persistent else "nonpersistent"
    varlen_str = "_varlen" if varlen else ""
    lse_str = "_lse" if with_lse else ""
    skip_str = "_skipsm" if enable_skip_softmax else ""
    sink_str = "_sink" if enable_sink else ""
    pdl_str = "_pdl" if use_pdl else ""
    ffi_str = "_tvmffi" if enable_tvm_ffi else ""
    return f"cute_dsl_fmha_{dtype_str}_h{head_dim}_{causal_str}_{persist_str}{varlen_str}{lse_str}{skip_str}{sink_str}{pdl_str}{ffi_str}"


# =============================================================================
# Loading: ExternalBinaryModule path
# =============================================================================


def _load_from_artifact(variant_name: str, gpu_arch: str, enable_tvm_ffi: bool = False):
    """Download .so from artifactory and load via ExternalBinaryModule.

    This is the production path used when cubins are published.

    Parameters
    ----------
    variant_name : str
        The kernel variant name (matches function_prefix used during export).
    gpu_arch : str
        GPU architecture string (e.g. 'sm_100a').
    enable_tvm_ffi : bool
        If False (default), load with CuTe native ABI.
        If True, load with TVM-FFI ABI (TODO: compile-side support pending).
    """
    from flashinfer.jit.cubin_loader import get_artifact
    from flashinfer.jit.env import FLASHINFER_CUBIN_DIR

    so_filename = f"{variant_name}.so"
    cpu_arch = _get_cpu_arch()
    artifact_path = os.path.join(
        DSL_FMHA_ARTIFACT_PATH, cpu_arch, gpu_arch, so_filename
    )

    checksums = _get_checksums(gpu_arch)
    sha256 = checksums.get(variant_name, "")
    if not sha256:
        raise RuntimeError(
            f"No checksum found for DSL FMHA variant '{variant_name}' on {gpu_arch}. "
            f"Available variants: {list(checksums.keys())}"
        )

    # get_artifact downloads to FLASHINFER_CUBIN_DIR / artifact_path and verifies sha256
    data = get_artifact(artifact_path, sha256)
    if not data:
        raise RuntimeError(f"Failed to download DSL FMHA artifact: {artifact_path}")

    local_path = FLASHINFER_CUBIN_DIR / artifact_path
    module = cute.runtime.load_module(str(local_path), enable_tvm_ffi=enable_tvm_ffi)
    return getattr(module, variant_name)


def _load_from_local(variant_name: str, local_dir: str, enable_tvm_ffi: bool = False):
    """Load .so or .o from a local directory (for development/testing).

    Set FLASHINFER_DSL_FMHA_LOCAL_DIR to use this path.

    Parameters
    ----------
    variant_name : str
        The kernel variant name (matches function_prefix used during export).
    local_dir : str
        Directory containing the compiled .so/.o files.
    enable_tvm_ffi : bool
        If False (default), load with CuTe native ABI.
        If True, load with TVM-FFI ABI (TODO: compile-side support pending).
    """
    # Try .so first, then .o
    so_path = os.path.join(local_dir, f"{variant_name}.so")
    o_path = os.path.join(local_dir, f"{variant_name}.o")
    if os.path.exists(so_path):
        load_path = so_path
    elif os.path.exists(o_path):
        load_path = o_path
    else:
        raise FileNotFoundError(
            f"DSL FMHA .so/.o not found at {local_dir}/{variant_name}.[so|o]. "
            f"Run compile_cute_dsl_fmha.py to generate it."
        )

    module = cute.runtime.load_module(load_path, enable_tvm_ffi=enable_tvm_ffi)
    return getattr(module, variant_name)


@functools.cache
def get_cute_dsl_fmha_kernel(
    gpu_arch: str,
    qk_dtype: torch.dtype,
    pv_dtype: torch.dtype,
    out_dtype: torch.dtype,
    head_dim: int,
    is_causal: bool,
    is_persistent: bool = True,
    enable_tvm_ffi: bool = False,
    varlen: bool = False,
    with_lse: bool = False,
    enable_skip_softmax: bool = False,
    enable_sink: bool = False,
    use_pdl: bool = False,
):
    """Get a compiled DSL FMHA kernel function.

    Checks local directory first (FLASHINFER_DSL_FMHA_LOCAL_DIR env var),
    then falls back to artifact download.

    Parameters
    ----------
    gpu_arch : str
        GPU architecture string (e.g. 'sm_100a'), used both for selecting the
        correct artifact and as part of the cache key.
    qk_dtype : torch.dtype
        Query/key data type (torch.float16, torch.bfloat16, or
        torch.float8_e4m3fn).
    pv_dtype : torch.dtype
        Value data type (torch.float16, torch.bfloat16, or
        torch.float8_e4m3fn).
    out_dtype : torch.dtype
        Output data type. Same as qk_dtype for non-mixed precision.
    head_dim : int
        Head dimension (e.g., 64, 128, 192). Note: 192 only supports FP8.
    is_causal : bool
        Whether to use causal masking.
    is_persistent : bool
        Whether to use persistent kernel mode.
    enable_tvm_ffi : bool
        If False (default), load with CuTe native ABI — kernel accepts
        cute Pointer/Tensor args (same calling convention as JIT mode).
        If True, load with TVM-FFI ABI — Pointer args accept data_ptr(),
        Tensor args accept torch.Tensor directly, stream uses env stream.
    enable_skip_softmax : bool
        If True, load kernel compiled with skip-softmax support.

    Returns
    -------
    callable
        The compiled kernel function.
    """
    variant_name = _get_variant_name(
        qk_dtype,
        pv_dtype,
        out_dtype,
        head_dim,
        is_causal,
        is_persistent,
        varlen=varlen,
        with_lse=with_lse,
        enable_tvm_ffi=enable_tvm_ffi,
        enable_skip_softmax=enable_skip_softmax,
        enable_sink=enable_sink,
        use_pdl=use_pdl,
    )

    # Check for local .so directory (development mode)
    local_dir = os.environ.get("FLASHINFER_DSL_FMHA_LOCAL_DIR")
    if local_dir:
        logger.info(
            f"Loading DSL FMHA kernel from local dir: {local_dir} (tvm_ffi={enable_tvm_ffi})"
        )
        return _load_from_local(variant_name, local_dir, enable_tvm_ffi=enable_tvm_ffi)

    # Production path: download from artifactory
    logger.info(
        f"Loading DSL FMHA kernel variant: {variant_name} (tvm_ffi={enable_tvm_ffi})"
    )
    return _load_from_artifact(variant_name, gpu_arch, enable_tvm_ffi=enable_tvm_ffi)


# =============================================================================
# PyTorch API wrapper
# =============================================================================


def cute_dsl_fmha_ragged_prefill(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    qo_indptr: torch.Tensor,
    kv_indptr: torch.Tensor,
    is_causal: bool = False,
    sm_scale: Optional[float] = None,
    window_left: int = -1,
    window_right: int = -1,
    lse: Optional[torch.Tensor] = None,
    attention_sinks: Optional[torch.Tensor] = None,
    scale_q: float = 1.0,
    scale_k: float = 1.0,
    scale_v: float = 1.0,
    scale_o: float = 1.0,
    enable_tvm_ffi: bool = True,
    max_qo_len: Optional[int] = None,
    max_kv_len: Optional[int] = None,
    kernel_fn=None,
    skip_softmax_threshold_scale_factor: Optional[float] = None,
    enable_pdl: bool = False,
) -> None:
    """Run DSL FMHA prefill kernel on ragged (variable-length) tensors.

    Note: The DSL FMHA kernel only supports per-tensor scalar scales, not
    per-head scale tensors.

    Parameters
    ----------
    q : torch.Tensor
        Query tensor, shape (total_q_tokens, H_q, D).
    k : torch.Tensor
        Key tensor, shape (total_kv_tokens, H_k, D).
    v : torch.Tensor
        Value tensor, shape (total_kv_tokens, H_k, D_v).
    o : torch.Tensor
        Output tensor, shape (total_q_tokens, H_q, D_v). Modified in-place.
        Must be 32-byte aligned (kernel uses 256-bit store instructions).
    qo_indptr : torch.Tensor
        Cumulative sequence lengths for Q/O, shape (batch_size + 1,).
        Same as cum_seqlen_q in DSL FMHA kernel.
    kv_indptr : torch.Tensor
        Cumulative sequence lengths for K/V, shape (batch_size + 1,).
        Same as cum_seqlen_k in DSL FMHA kernel.
    is_causal : bool
        Whether to apply causal masking.
    sm_scale : float, optional
        Softmax scale factor. Defaults to 1/sqrt(D).
    window_left : int
        Left sliding window size. -1 means no window.
    window_right : int
        Right sliding window size. -1 means no window. 0 for causal.
    lse : torch.Tensor, optional
        Log-sum-exp output tensor. None to skip.
    attention_sinks : torch.Tensor, optional
        Attention sink tensor, shape (H_q,) Float32. None to disable.
    scale_q : float
        Per-tensor scale for query (FP8 calibration). Default 1.0.
    scale_k : float
        Per-tensor scale for key (FP8 calibration). Default 1.0.
    scale_v : float
        Per-tensor scale for value (FP8 calibration). Default 1.0.
    scale_o : float
        Per-tensor scale for output (FP8 calibration). Default 1.0.
    enable_tvm_ffi : bool
        If True, use TVM-FFI ABI (pass data_ptr() for Pointer args, torch.Tensor
        for Tensor args, no explicit stream). Default False (CuTe native ABI).
    max_qo_len : int, optional
        Maximum query sequence length. Computed from qo_indptr if not provided.
        Pass this from plan() to avoid D2H copy during CUDA graph capture.
    max_kv_len : int, optional
        Maximum KV sequence length. Computed from kv_indptr if not provided.
    skip_softmax_threshold_scale_factor : float, optional
        Threshold scale factor for skip-softmax sparsity (https://arxiv.org/abs/2512.12087).
        The actual threshold = scale_factor / max_kv_len, then converted to log2 domain.
        None or 0 disables skip-softmax.
    """
    total_q, H_q, D = q.shape
    total_kv, H_k, _ = k.shape
    D_v = v.shape[-1]

    batch_size = len(qo_indptr) - 1

    use_skip_softmax = (
        skip_softmax_threshold_scale_factor is not None
        and skip_softmax_threshold_scale_factor > 0
    )

    if kernel_fn is None:
        gpu_arch = _get_gpu_arch(q.device)
        # Q and K feed the same GEMM operand dtype: neither the prebuilt
        # artifacts nor the JIT compile (single qk_dtype) support q != k.
        if q.dtype != k.dtype:
            raise ValueError(
                f"cute-dsl FMHA requires q.dtype == k.dtype, got {q.dtype} and {k.dtype}"
            )
        try:
            if window_left != -1 or window_right != -1:
                # The exported artifact matrix has no window axis: every
                # prebuilt variant is traced with window_size_left=None (and
                # window_size_right present only for causal), so a windowed
                # call into one fails the FFI signature check.  The causal
                # artifact lookup would falsely succeed here — force JIT.
                raise RuntimeError("windowed DSL FMHA variants are not exported")
            kernel_fn = get_cute_dsl_fmha_kernel(
                gpu_arch,
                q.dtype,
                v.dtype,
                o.dtype,
                D,
                is_causal,
                is_persistent=False,  # varlen always uses non-persistent
                varlen=True,
                enable_tvm_ffi=enable_tvm_ffi,
                with_lse=lse is not None,
                enable_skip_softmax=use_skip_softmax,
                enable_sink=attention_sinks is not None,
                use_pdl=enable_pdl,
            )
        except (RuntimeError, FileNotFoundError) as e:
            # Cubin variant not available / repo unreachable.
            # JIT-compile the trtllm kernel instead.
            logger.info(f"DSL FMHA cubin unavailable ({e}); JIT-compiling the kernel.")
            from flashinfer.cute_dsl.attention.fmha.compile import (
                compile_cute_dsl_fmha_kernel,
            )

            kernel_fn = compile_cute_dsl_fmha_kernel(
                q.dtype,
                v.dtype,
                o.dtype,
                H_q,
                H_k,
                D,
                D_v,
                is_causal,
                lse is not None,
                attention_sinks is not None,
                use_skip_softmax,
                enable_pdl,
                q.device,
                has_window_left=window_left != -1,
                has_window_right=is_causal or window_right != -1,
            )
            # JIT kernels are compiled with --enable-tvm-ffi; the CuTe-native branch
            # would not accept them.
            enable_tvm_ffi = True

    # Compute scale factors
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(D)
    log2_e = math.log2(math.exp(1.0))
    scale_softmax = scale_q * scale_k * sm_scale
    scale_softmax_log2 = scale_softmax * log2_e
    scale_output = scale_v / scale_o

    # Max seq lengths for problem_size (prefer pre-computed values for CUDA graph compat)
    if max_qo_len is None:
        max_s_q = int((qo_indptr.cpu()[1:] - qo_indptr.cpu()[:-1]).max().item())
    else:
        max_s_q = max_qo_len
    if max_kv_len is None:
        max_s_k = int((kv_indptr.cpu()[1:] - kv_indptr.cpu()[:-1]).max().item())
    else:
        max_s_k = max_kv_len

    # problem_size: (B, max_s_q, s_lse, max_s_k, H_q, H_k, D, D_v)
    s_lse = total_q  # for variable length, s_lse = total tokens
    problem_size = (batch_size, max_s_q, s_lse, max_s_k, H_q, H_k, D, D_v)

    # Skip-softmax threshold (convert scale_factor to log2 domain)
    skip_threshold_log2 = None
    if use_skip_softmax:
        threshold = skip_softmax_threshold_scale_factor / max_s_k
        skip_threshold_log2 = Float32(math.log2(threshold))

    # Window sizes
    ws_left = None if window_left == -1 else Int32(window_left)
    ws_right = None if window_right == -1 else Int32(window_right)
    if is_causal and ws_right is None:
        ws_right = Int32(0)

    # Reshape to 5D matching kernel docstring:
    #   q/o: (b=1, total, h_k, h_r, d/dv)
    #   k/v: (b=1, total, h_k, 1, d/dv)
    h_r = H_q // H_k
    q_5d = q.view(1, total_q, H_k, h_r, D)
    k_5d = k.view(1, total_kv, H_k, 1, D)
    v_5d = v.view(1, total_kv, H_k, 1, D_v)
    assert o.data_ptr() % 32 == 0, (
        "o must be 32-byte aligned (kernel uses 256-bit store instructions)"
    )
    o_5d = o.view(1, total_q, H_k, h_r, D_v)
    # LSE: (1, total_q, h_k, h_r) — 4D row-major.
    lse_4d = lse.view(1, total_q, H_k, h_r) if lse is not None else None

    if enable_tvm_ffi:
        # TVM-FFI: pass torch.Tensor directly, no explicit stream (env stream).
        kernel_fn(
            q_5d,
            k_5d,
            v_5d,
            o_5d,
            problem_size,
            qo_indptr.to(torch.int32),  # cum_seqlen_q
            kv_indptr.to(torch.int32),  # cum_seqlen_k
            lse_4d,
            attention_sinks,
            Float32(scale_softmax_log2),
            Float32(scale_softmax),
            Float32(scale_output),
            skip_threshold_log2,
            ws_left,
            ws_right,
            None,  # skip_softmax_count
            None,  # total_softmax_count
            enable_pdl,
        )
    else:
        # CuTe native ABI: convert to cute tensors and pass with explicit stream.
        is_fp8_qk = q.dtype == torch.float8_e4m3fn
        is_fp8_pv = v.dtype == torch.float8_e4m3fn
        is_fp8_out = o.dtype == torch.float8_e4m3fn
        if is_fp8_qk:
            q_cute = from_dlpack(
                q_5d.view(torch.int8), assumed_align=16
            ).mark_layout_dynamic(leading_dim=4)
            q_cute.element_type = cutlass.Float8E4M3FN
            k_cute = from_dlpack(
                k_5d.view(torch.int8), assumed_align=16
            ).mark_layout_dynamic(leading_dim=4)
            k_cute.element_type = cutlass.Float8E4M3FN
        else:
            q_cute = from_dlpack(q_5d, assumed_align=16).mark_layout_dynamic(
                leading_dim=4
            )
            k_cute = from_dlpack(k_5d, assumed_align=16).mark_layout_dynamic(
                leading_dim=4
            )
        if is_fp8_pv:
            v_cute = from_dlpack(
                v_5d.view(torch.int8), assumed_align=16
            ).mark_layout_dynamic(leading_dim=4)
            v_cute.element_type = cutlass.Float8E4M3FN
        else:
            v_cute = from_dlpack(v_5d, assumed_align=16).mark_layout_dynamic(
                leading_dim=4
            )
        if is_fp8_out:
            o_cute = from_dlpack(
                o_5d.view(torch.int8), assumed_align=32
            ).mark_layout_dynamic(leading_dim=4)
            o_cute.element_type = cutlass.Float8E4M3FN
        else:
            o_cute = from_dlpack(o_5d, assumed_align=32).mark_layout_dynamic(
                leading_dim=4
            )

        cum_seqlen_q_cute = from_dlpack(
            qo_indptr.to(torch.int32), assumed_align=16
        ).mark_layout_dynamic(leading_dim=0)
        cum_seqlen_k_cute = from_dlpack(
            kv_indptr.to(torch.int32), assumed_align=16
        ).mark_layout_dynamic(leading_dim=0)

        lse_cute = None
        if lse_4d is not None:
            lse_cute = from_dlpack(lse_4d, assumed_align=16).mark_layout_dynamic(
                leading_dim=3
            )

        sink_cute = None
        if attention_sinks is not None:
            sink_cute = from_dlpack(
                attention_sinks, assumed_align=16
            ).mark_layout_dynamic(leading_dim=0)

        stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream)

        kernel_fn(
            q_cute,
            k_cute,
            v_cute,
            o_cute,
            problem_size,
            cum_seqlen_q_cute,
            cum_seqlen_k_cute,
            lse_cute,
            sink_cute,
            Float32(scale_softmax_log2),
            Float32(scale_softmax),
            Float32(scale_output),
            skip_threshold_log2,
            ws_left,
            ws_right,
            None,  # skip_softmax_count
            None,  # total_softmax_count
            stream,
            enable_pdl,
        )
