import os

import torch
import torch.nn.functional as F
from loguru import logger

from text_embeddings_server.utils.device import is_hpu, is_rocm, use_ipex

if os.getenv("USE_FLASH_ATTENTION", "").lower() == "false":
    raise ImportError("`USE_FLASH_ATTENTION` is false.")

HAS_FLASH_ATTN = False
HAS_FLASH_ATTN_V2 = False

is_hpu = is_hpu()
is_rocm = is_rocm()
use_ipex = use_ipex()

if use_ipex or is_hpu:
    HAS_FLASH_ATTN_V2 = True
elif is_rocm:
    HAS_FLASH_ATTN_V2 = True

    try:
        from flash_attn import flash_attn_varlen_func as _fa_varlen_func
        ROCM_HAS_FA_VARLEN = True
        logger.info("ROCm: flash_attn found, using flash_attn_varlen_func")
    except ImportError:
        ROCM_HAS_FA_VARLEN = False
        logger.warning(
            "ROCm: flash_attn not found, falling back to SDPA. "
            "Install flash-attn for better performance."
        )
else:
    if not torch.cuda.is_available():
        raise ImportError("CUDA is not available")

    major, minor = torch.cuda.get_device_capability()
    is_sm75 = major == 7 and minor == 5
    is_sm8x = major == 8 and minor >= 0
    is_sm90 = major == 9 and minor == 0

    try:
        try:
            import flash_attn_2_cuda
        except ImportError:
            raise ImportError(
                "Flash Attention V2 is not installed.\n"
                "Use the official Docker image (ghcr.io/huggingface/text-embeddings-inference:cuda-latest) "
                "or install flash attention v2 with `cd server && make install install-flash-attention-v2`"
            )
        if not (is_sm8x or is_sm90):
            raise ImportError(
                f"GPU with CUDA capability {major} {minor} is not supported for "
                "Flash Attention V2"
            )
        HAS_FLASH_ATTN_V2 = True
    except ImportError as e:
        try:
            import flash_attn_cuda
        except ImportError:
            raise ImportError(
                "Flash Attention is not installed.\n"
                "Use the official Docker image (ghcr.io/huggingface/text-embeddings-inference:cuda-latest) "
                "or install flash attention with `cd server && make install install-flash-attention`"
            ) from e

        if not (is_sm75 or is_sm8x or is_sm90):
            raise ImportError(
                f"GPU with CUDA capability {major} {minor} is not supported"
            ) from e
        logger.warning(f"Unable to use Flash Attention V2: {e}")
        HAS_FLASH_ATTN = True


def rocm_attn(
    q,
    k,
    v,
    out,
    attn_mask,
    softmax_scale,
    is_causal=False,
):
    # q/k/v arrive as [bs, seq_len, heads, head_dim] (PaddedBatch layout)
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    result = F.scaled_dot_product_attention(
        q,
        k,
        v,
        attn_mask=attn_mask,
        dropout_p=0.0,
        is_causal=is_causal,
        scale=softmax_scale,
    )
    result = result.transpose(1, 2)
    out.copy_(result)
    return out


def hpu_attn(
    q,
    k,
    v,
    out,
    attn_mask,
    seqlen_q,
    seqlen_k,
    max_seqlen_q,
    max_seqlen_k,
    softmax_scale,
    is_causal=False,
):
    from habana_frameworks.torch.hpex.kernels import FusedSDPA

    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    if is_causal:
        attn_mask = None

    out_ = FusedSDPA.apply(
        q, k, v, attn_mask, 0.0, is_causal, softmax_scale, "fast", False
    )
    out_ = out_.transpose(1, 2)
    out.copy_(out_)
    return out


def attention(
    q, k, v, out, cu_seqlens, max_s, softmax_scale, is_causal=False, attn_mask=None
):
    if use_ipex and HAS_FLASH_ATTN_V2:
        import intel_extension_for_pytorch as ipex

        return ipex.llm.functional.varlen_attention(
            q.contiguous() if q.device.type == "xpu" else q,
            k.contiguous() if k.device.type == "xpu" else k,
            v.contiguous() if v.device.type == "xpu" else v,
            out,
            cu_seqlens,
            cu_seqlens,
            None,
            max_s,
            max_s,
            0,
            softmax_scale,
            zero_tensors=False,
            is_causal=False,
            return_softmax=False,
            gen_=None,
        )

    if is_rocm:
        if ROCM_HAS_FA_VARLEN:
            result = _fa_varlen_func(
                q, k, v,
                cu_seqlens, cu_seqlens,
                max_s, max_s,
                causal=is_causal,
                softmax_scale=softmax_scale,
            )
            out.copy_(result)
            return out
        else:
            return rocm_attn(q, k, v, out, attn_mask, softmax_scale, is_causal)

    if is_hpu:
        return hpu_attn(
            q,
            k,
            v,
            out,
            attn_mask,
            cu_seqlens,
            cu_seqlens,
            max_s,
            max_s,
            softmax_scale,
            is_causal,
        )

    if HAS_FLASH_ATTN_V2:
        return flash_attn_2_cuda.varlen_fwd(
            q,
            k,
            v,
            out,
            cu_seqlens,
            cu_seqlens,
            max_s,
            max_s,
            0.0,
            softmax_scale,
            False,
            is_causal,
            -1,
            -1,
            False,
            None,
        )

    if HAS_FLASH_ATTN:
        return flash_attn_cuda.fwd(
            q,
            k,
            v,
            out,
            cu_seqlens,
            cu_seqlens,
            max_s,
            max_s,
            0.0,
            softmax_scale,
            False,
            is_causal,
            False,
            0,
            None,
        )

    raise NotImplementedError("flash attention is not installed")