# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import torch
from typing import Optional
import os
import warnings

# Global state for lazy initialization
_SAGEATTN_AVAILABLE = None
_FLASH_ATTN_3_AVAILABLE = None
_FLASH_ATTN_2_AVAILABLE = None
_sageattn_func = None
_flash_attn_func = None
_flash_attn_interface = None
_flash_attn = None


def _init_sageattention():
    """Lazy initialization for SageAttention."""
    global _SAGEATTN_AVAILABLE, _sageattn_func

    if _SAGEATTN_AVAILABLE is not None:
        return _SAGEATTN_AVAILABLE

    _SAGEATTN_AVAILABLE = False
    try:
        if os.getenv("DISABLE_SAGEATTENTION", "0") != "0":
            raise Exception("DISABLE_SAGEATTENTION is set")

        from sageattention import sageattn

        @torch.library.custom_op(
            "mylib::sageattn", mutates_args={"q", "k", "v"}, device_types="cuda"
        )
        def sageattn_func(
            q: torch.Tensor,
            k: torch.Tensor,
            v: torch.Tensor,
            attn_mask: Optional[torch.Tensor] = None,
            dropout_p: float = 0,
            is_causal: bool = False,
        ) -> torch.Tensor:
            return sageattn(
                q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal
            )

        @sageattn_func.register_fake
        def _sageattn_fake(q, k, v, attn_mask=None, dropout_p=0, is_causal=False):
            return torch.empty(*q.shape, device=q.device, dtype=q.dtype)

        print("SageAttention loaded successfully")
        _sageattn_func = sageattn_func
        _SAGEATTN_AVAILABLE = True

    except Exception as e:
        print(f"Warning: Could not load sageattention: {str(e)}")
        if isinstance(e, ModuleNotFoundError):
            print("sageattention package is not installed")
        elif isinstance(e, ImportError) and "DLL" in str(e):
            print("sageattention DLL loading error")
        _sageattn_func = None

    return _SAGEATTN_AVAILABLE


def _is_hopper_gpu():
    """Check if the current GPU is a Hopper architecture."""
    if not torch.cuda.is_available():
        return False
    device_name = torch.cuda.get_device_name(0).lower()
    return "h100" in device_name or "hopper" in device_name


def _init_flash_attention_3():
    """Lazy initialization for Flash Attention 3."""
    global _FLASH_ATTN_3_AVAILABLE, _flash_attn_func, _flash_attn_interface

    if _FLASH_ATTN_3_AVAILABLE is not None:
        return _FLASH_ATTN_3_AVAILABLE

    _FLASH_ATTN_3_AVAILABLE = False
    try:
        from flash_attn import flash_attn_func
        import flash_attn_interface

        # Always set the function reference if flash_attn is available
        _flash_attn_func = flash_attn_func
        _flash_attn_interface = flash_attn_interface
        # FA3 optimizations only available on Hopper GPUs
        _FLASH_ATTN_3_AVAILABLE = _is_hopper_gpu()
    except ModuleNotFoundError:
        _FLASH_ATTN_3_AVAILABLE = False
        _flash_attn_func = None
        _flash_attn_interface = None

    return _FLASH_ATTN_3_AVAILABLE


def _init_flash_attention_2():
    """Lazy initialization for Flash Attention 2."""
    global _FLASH_ATTN_2_AVAILABLE, _flash_attn

    if _FLASH_ATTN_2_AVAILABLE is not None:
        return _FLASH_ATTN_2_AVAILABLE

    _FLASH_ATTN_2_AVAILABLE = False
    try:
        import flash_attn

        _flash_attn = flash_attn
        _FLASH_ATTN_2_AVAILABLE = True
    except ModuleNotFoundError:
        _FLASH_ATTN_2_AVAILABLE = False

    return _FLASH_ATTN_2_AVAILABLE

__all__ = ["flash_attention", "attention"]


# Compatibility getters for external code
def sageattn_func():
    """Getter for sageattn_func - initializes if needed."""
    _init_sageattention()
    return _sageattn_func


def SAGEATTN_AVAILABLE():
    """Getter for SAGEATTN_AVAILABLE - initializes if needed."""
    return _init_sageattention()


def flash_attention(
    q,
    k,
    v,
    q_lens=None,
    k_lens=None,
    dropout_p=0.0,
    softmax_scale=None,
    q_scale=None,
    causal=False,
    window_size=(-1, -1),
    deterministic=False,
    dtype=torch.bfloat16,
    version=None,
):
    """
    q:              [B, Lq, Nq, C1].
    k:              [B, Lk, Nk, C1].
    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
    q_lens:         [B].
    k_lens:         [B].
    dropout_p:      float. Dropout probability.
    softmax_scale:  float. The scaling of QK^T before applying softmax.
    causal:         bool. Whether to apply causal attention mask.
    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
    deterministic:  bool. If True, slightly slower and uses more memory.
    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
    """
    # Initialize flash attention modules
    flash_attn_3_available = _init_flash_attention_3()
    flash_attn_2_available = _init_flash_attention_2()

    # Early fallback for simple cases when advanced features aren't needed
    # Only use this path if flash_attn is available but we're not using FA3 features
    if not flash_attn_3_available and _flash_attn_func is not None and q_lens is None and k_lens is None:
        return _flash_attn_func(
            q,
            k,
            v,
        )

    half_dtypes = (torch.float16, torch.bfloat16)
    assert dtype in half_dtypes
    assert q.device.type == "cuda" and q.size(-1) <= 256

    # params
    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype

    def half(x):
        return x if x.dtype in half_dtypes else x.to(dtype)

    # preprocess query
    if q_lens is None:
        q = half(q.flatten(0, 1))
        q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(
            device=q.device, non_blocking=True
        )
    else:
        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))

    # preprocess key, value
    if k_lens is None:
        k = half(k.flatten(0, 1))
        v = half(v.flatten(0, 1))
        k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(
            device=k.device, non_blocking=True
        )
    else:
        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))

    q = q.to(v.dtype)
    k = k.to(v.dtype)

    if q_scale is not None:
        q = q * q_scale

    if version is not None and version == 3 and not flash_attn_3_available:
        warnings.warn(
            "Flash attention 3 is not available, use flash attention 2 instead."
        )

    # apply attention
    if (version is None or version == 3) and flash_attn_3_available:
        # Note: dropout_p, window_size are not supported in FA3 now.
        x = _flash_attn_interface.flash_attn_varlen_func(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
            .cumsum(0, dtype=torch.int32)
            .to(q.device, non_blocking=True),
            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
            .cumsum(0, dtype=torch.int32)
            .to(q.device, non_blocking=True),
            max_seqlen_q=lq,
            max_seqlen_k=lk,
            softmax_scale=softmax_scale,
            causal=causal,
            deterministic=deterministic,
        ).unflatten(0, (b, lq))
    else:
        assert flash_attn_2_available
        x = _flash_attn.flash_attn_varlen_func(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
            .cumsum(0, dtype=torch.int32)
            .to(q.device, non_blocking=True),
            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
            .cumsum(0, dtype=torch.int32)
            .to(q.device, non_blocking=True),
            max_seqlen_q=lq,
            max_seqlen_k=lk,
            dropout_p=dropout_p,
            softmax_scale=softmax_scale,
            causal=causal,
            window_size=window_size,
            deterministic=deterministic,
        ).unflatten(0, (b, lq))

    # output
    return x.type(out_dtype)


def attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q_lens=None,
    k_lens=None,
    dropout_p=0.0,
    softmax_scale=None,
    q_scale=None,
    causal=False,
    window_size=(-1, -1),
    deterministic=False,
    dtype=torch.bfloat16,
    fa_version=None,
    # og_dtype=torch.bfloat16,
):
    # Initialize attention modules
    sageattn_available = _init_sageattention()
    flash_attn_2_available = _init_flash_attention_2()
    flash_attn_3_available = _init_flash_attention_3()

    if sageattn_available:
        # print("Using sageattention")
        attn_mask = None

        og_dtype = q.dtype
        q = q.transpose(1, 2).to(dtype)
        k = k.transpose(1, 2).to(dtype)
        v = v.transpose(1, 2).to(dtype)

        out = _sageattn_func(
            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
        )

        out = out.transpose(1, 2).contiguous().to(og_dtype)
        return out

    elif flash_attn_2_available or flash_attn_3_available:
        return flash_attention(
            q=q,
            k=k,
            v=v,
            q_lens=q_lens,
            k_lens=k_lens,
            dropout_p=dropout_p,
            softmax_scale=softmax_scale,
            q_scale=q_scale,
            causal=causal,
            window_size=window_size,
            deterministic=deterministic,
            dtype=dtype,
            version=fa_version,
        )
    else:
        if q_lens is not None or k_lens is not None:
            warnings.warn(
                "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance."
            )
        attn_mask = None

        q = q.transpose(1, 2).to(dtype)
        k = k.transpose(1, 2).to(dtype)
        v = v.transpose(1, 2).to(dtype)

        out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
        )

        out = out.transpose(1, 2).contiguous()
        return out