# SPDX-License-Identifier: Apache-2.0
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py
"""Code inside this file can safely assume cuda platform, e.g. importing
pynvml. However, it should not initialize cuda context.
"""

import ctypes
import os
from collections.abc import Callable
from functools import lru_cache, wraps
from typing import TypeVar

import torch
from typing_extensions import ParamSpec

import fastvideo.envs as envs
from fastvideo.logger import init_logger
from fastvideo.platforms.interface import (AttentionBackendEnum, DeviceCapability, Platform, PlatformEnum)
from fastvideo.utils import import_pynvml

logger = init_logger(__name__)

_P = ParamSpec("_P")
_R = TypeVar("_R")

pynvml = import_pynvml()  # type: ignore[no-untyped-call]

_CUDA_SUCCESS = 0
# Stable value from CUDA's public CUdevice_attribute enum.
_CU_DEVICE_ATTRIBUTE_INTEGRATED = 18
_CUDA_DRIVER_LIBRARY = "nvcuda.dll" if os.name == "nt" else "libcuda.so.1"


def _cuda_driver_device_is_integrated(device_id: int) -> bool:
    """Query ``CU_DEVICE_ATTRIBUTE_INTEGRATED`` without creating a context.

    ``cuInit`` loads the CUDA driver, while these device-management calls only
    inspect the logical device ordinal. They neither create nor retain a CUDA
    context. Driver initialization is not pre-fork safe, so this probe must
    remain worker-local after process creation and device binding. Using the
    driver ordinal preserves CUDA_VISIBLE_DEVICES ordering, including UUID and
    MIG selectors.
    """
    try:
        driver = ctypes.CDLL(_CUDA_DRIVER_LIBRARY)
        driver.cuInit.argtypes = [ctypes.c_uint]
        driver.cuInit.restype = ctypes.c_int
        driver.cuDeviceGet.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int]
        driver.cuDeviceGet.restype = ctypes.c_int
        driver.cuDeviceGetAttribute.argtypes = [
            ctypes.POINTER(ctypes.c_int),
            ctypes.c_int,
            ctypes.c_int,
        ]
        driver.cuDeviceGetAttribute.restype = ctypes.c_int

        if driver.cuInit(0) != _CUDA_SUCCESS:
            return False

        device = ctypes.c_int()
        if driver.cuDeviceGet(ctypes.byref(device), device_id) != _CUDA_SUCCESS:
            return False

        is_integrated = ctypes.c_int()
        if driver.cuDeviceGetAttribute(
                ctypes.byref(is_integrated),
                _CU_DEVICE_ATTRIBUTE_INTEGRATED,
                device,
        ) != _CUDA_SUCCESS:
            return False
        return bool(is_integrated.value)
    except Exception:
        # Missing/incompatible driver libraries and unavailable devices must
        # preserve the established discrete-memory offload policy.
        return False


# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
# see https://github.com/huggingface/diffusers/issues/9704 for details
torch.backends.cuda.enable_cudnn_sdp(False)


def device_id_to_physical_device_id(device_id: int) -> int:
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
        if device_ids == [""]:
            msg = ("CUDA_VISIBLE_DEVICES is set to empty string, which means"
                   " GPU support is disabled. If you are using ray, please unset"
                   " the environment variable `CUDA_VISIBLE_DEVICES` inside the"
                   " worker/actor. "
                   "Check https://github.com/vllm-project/vllm/issues/8402 for"
                   " more information.")
            raise RuntimeError(msg)
        physical_device_id = device_ids[device_id]
        return int(physical_device_id)
    else:
        return device_id


def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:

    @wraps(fn)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
        pynvml.nvmlInit()
        try:
            return fn(*args, **kwargs)
        finally:
            pynvml.nvmlShutdown()

    return wrapper


class CudaPlatformBase(Platform):
    _enum = PlatformEnum.CUDA
    device_name: str = "cuda"
    device_type: str = "cuda"
    dispatch_key: str = "CUDA"
    ray_device_key: str = "GPU"
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"

    @classmethod
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
        raise NotImplementedError

    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        raise NotImplementedError

    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        raise NotImplementedError

    @classmethod
    def has_unified_memory(cls, device_id: int = 0) -> bool:
        # This is cudaDeviceProp::integrated's driver-level source of truth. It
        # is true on parts such as GB10 and Jetson whose GPU reads host memory.
        return _cuda_driver_device_is_integrated(device_id)

    @classmethod
    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
        if enforce_eager:
            logger.warning("To see benefits of async output processing, enable CUDA "
                           "graph. Since, enforce-eager is enabled, async output "
                           "processor cannot be used")
            return False
        return True

    @classmethod
    def is_full_nvlink(cls, device_ids: list[int]) -> bool:
        raise NotImplementedError

    @classmethod
    def log_warnings(cls) -> None:
        pass

    @classmethod
    def get_current_memory_usage(cls, device: torch.types.Device | None = None) -> float:
        torch.cuda.reset_peak_memory_stats(device)
        return float(torch.cuda.max_memory_allocated(device))

    @classmethod
    def get_torch_device(cls):
        """
        Return torch.cuda
        """
        return torch.cuda

    @classmethod
    def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int,
                             dtype: torch.dtype) -> str:
        # TODO(will): maybe come up with a more general interface for local attention
        # if distributed is False, we always try to use Flash attn

        logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND)
        logger.info("Selected backend: %s", selected_backend)
        if selected_backend == AttentionBackendEnum.SAGE_ATTN:
            try:
                from sageattention import sageattn  # noqa: F401

                from fastvideo.attention.backends.sage_attn import (  # noqa: F401
                    SageAttentionBackend)
                logger.info("Using Sage Attention backend.")

                return "fastvideo.attention.backends.sage_attn.SageAttentionBackend"
            except ImportError as e:
                logger.info(e)
                logger.info("Sage Attention backend is not installed. Fall back to Flash Attention.")
        elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE:
            try:
                from sageattn3 import sageattn3_blackwell  # noqa: F401

                from fastvideo.attention.backends.sage_attn3 import (  # noqa: F401
                    SageAttention3Backend)
                logger.info("Using Sage Attention 3 backend.")

                return "fastvideo.attention.backends.sage_attn3.SageAttention3Backend"
            except ImportError as e:
                logger.info(e)
                logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
        elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
            from fastvideo.attention.backends.attn_qat_infer import (  # noqa: F401
                AttnQatInferBackend, attn_qat_infer_receipt, is_attn_qat_infer_available)
            if is_attn_qat_infer_available():
                logger.info("Using Attn-QAT inference backend (%s).", attn_qat_infer_receipt())
                return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
            raise ImportError(
                f"ATTN_QAT_INFER selected but the inference kernel is not usable ({attn_qat_infer_receipt()}). "
                "Silent fallback would run plain FlashAttention while the caller believes it is measuring "
                "FP4-QAT attention — an A/B comparison would silently benchmark bf16 against bf16; "
                "refusing to proceed. Build the fastvideo-kernel attn_qat_infer target for this arch "
                "or pick a different FASTVIDEO_ATTENTION_BACKEND.")
        elif selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN:
            from fastvideo.attention.backends.attn_qat_train import (  # noqa: F401
                AttnQatTrainBackend, is_attn_qat_train_available)
            if is_attn_qat_train_available():
                logger.info("Using Attn-QAT training (fake-quantized attention) backend.")
                return "fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend"
            raise ImportError(
                "ATTN_QAT_TRAIN selected but fastvideo_kernel.triton_kernels.attn_qat_train is not built. "
                "Silent fallback would produce a non-QAT training run; refusing to proceed. "
                "Install the training kernel or pick a different FASTVIDEO_ATTENTION_BACKEND.")
        elif selected_backend == AttentionBackendEnum.NABLA_ATTN:
            from fastvideo.attention.backends.nabla import CAN_USE_FLEX_ATTN
            if CAN_USE_FLEX_ATTN:
                logger.info("Using NABLA block-sparse flex-attention backend.")
                return "fastvideo.attention.backends.nabla.NablaAttentionBackend"
            raise ImportError("NABLA_ATTN selected but torch.nn.attention.flex_attention is unavailable in this "
                              "PyTorch build. Silent fallback to dense attention would be orders of magnitude "
                              "slower and diverge from the reference; upgrade PyTorch or pick a different backend.")
        elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
            try:
                from fastvideo_kernel import video_sparse_attn  # noqa: F401

                from fastvideo.attention.backends.video_sparse_attn import (  # noqa: F401
                    VideoSparseAttentionBackend)
                logger.info("Using Video Sparse Attention backend.")

                return "fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionBackend"
            except ImportError as e:
                logger.error("Failed to import Video Sparse Attention backend: %s", str(e))
                raise ImportError("The Video Sparse Attention backend is not installed. "
                                  "To install it, please follow the instructions at: "
                                  "https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation ") from e
        elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3:
            try:
                from fastvideo_kernel.block_sparse_attn_256 import (  # noqa: F401
                    block_sparse_attn_256_bshd)

                logger.info("Using MiniMax-H3 Video Sparse Attention backend.")

                return "fastvideo.attention.backends.video_sparse_attn_h3.MiniMaxH3VSABackend"
            except ImportError as e:
                logger.error("Failed to import H3 Video Sparse Attention backend: %s", str(e))
                raise ImportError("VIDEO_SPARSE_ATTN_H3 selected but fastvideo_kernel's block-sparse "
                                  "kernels are unavailable. Install fastvideo-kernel or pick a different "
                                  "FASTVIDEO_ATTENTION_BACKEND.") from e
        elif selected_backend == AttentionBackendEnum.BSA_ATTN:
            try:
                from fastvideo.attention.backends.bsa_attn import (  # noqa: F401
                    BSAAttentionBackend)
                logger.info("Using BSA Attention backend.")

                return "fastvideo.attention.backends.bsa_attn.BSAAttentionBackend"
            except ImportError as e:
                logger.error("Failed to import BSA Attention backend: %s", str(e))
                raise ImportError("The BSA Attention backend failed to import.") from e
        elif selected_backend == AttentionBackendEnum.VMOBA_ATTN:
            try:
                from fastvideo_kernel import moba_attn_varlen  # noqa: F401
                from fastvideo.attention.backends.vmoba import (  # noqa: F401
                    VMOBAAttentionBackend)
                logger.info("Using Video MOBA Attention backend.")

                return "fastvideo.attention.backends.vmoba.VMOBAAttentionBackend"
            except ImportError as e:
                logger.error("Failed to import Video MoBA Attention backend: %s", str(e))
                raise ImportError("Video MoBA Attention backend is not installed. ") from e
        elif selected_backend == AttentionBackendEnum.SLA_ATTN:
            try:
                from fastvideo.attention.backends.sla import (  # noqa: F401
                    SLAAttentionBackend)
                logger.info("Using SLA (Sparse-Linear Attention) backend.")

                return "fastvideo.attention.backends.sla.SLAAttentionBackend"
            except ImportError as e:
                logger.error("Failed to import SLA Attention backend: %s", str(e))
                raise ImportError("SLA Attention backend is not available. ") from e
        elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN:
            try:
                from fastvideo.attention.backends.sla import (  # noqa: F401
                    SageSLAAttentionBackend)
                logger.info("Using SageSLA (Quantized Sparse-Linear Attention) backend.")

                return "fastvideo.attention.backends.sla.SageSLAAttentionBackend"
            except ImportError as e:
                logger.error("Failed to import SageSLA Attention backend: %s", str(e))
                raise ImportError("SageSLA Attention backend requires spas_sage_attn. "
                                  "Install with: uv pip install git+https://github.com/thu-ml/SpargeAttn.git") from e
        elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
            logger.info("Using Torch SDPA backend.")
            return "fastvideo.attention.backends.sdpa.SDPABackend"
        elif selected_backend == AttentionBackendEnum.FLASH_ATTN or selected_backend is None:
            pass
        elif selected_backend:
            raise ValueError(f"Invalid attention backend for {cls.device_name}")

        target_backend = AttentionBackendEnum.FLASH_ATTN
        if not cls.has_device_capability(80):
            logger.info("Cannot use FlashAttention-2 backend for Volta and Turing "
                        "GPUs.")
            target_backend = AttentionBackendEnum.TORCH_SDPA
        elif dtype not in (torch.float16, torch.bfloat16):
            logger.info("Cannot use FlashAttention-2 backend for dtype other than "
                        "torch.float16 or torch.bfloat16.")
            target_backend = AttentionBackendEnum.TORCH_SDPA

        # FlashAttn is valid for the model, checking if the package is
        # installed.
        if target_backend == AttentionBackendEnum.FLASH_ATTN:
            try:
                import flash_attn  # noqa: F401

                from fastvideo.attention.backends.flash_attn import (  # noqa: F401
                    FlashAttentionBackend)

                supported_sizes = \
                    FlashAttentionBackend.get_supported_head_sizes()
                if head_size not in supported_sizes:
                    logger.info("Cannot use FlashAttention-2 backend for head size %d.", head_size)
                    target_backend = AttentionBackendEnum.TORCH_SDPA
            except ImportError:
                logger.info("Cannot use FlashAttention-2 backend because the "
                            "flash_attn package is not found. "
                            "Make sure that flash_attn was built and installed "
                            "(on by default).")
                target_backend = AttentionBackendEnum.TORCH_SDPA

        if target_backend == AttentionBackendEnum.TORCH_SDPA:
            logger.info("Using Torch SDPA backend.")

            return "fastvideo.attention.backends.sdpa.SDPABackend"

        logger.info("Using Flash Attention backend.")

        return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend"

    @classmethod
    def get_device_communicator_cls(cls) -> str:
        return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa


# NVML utils
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
# all the related functions work on real physical device ids.
# the major benefit of using NVML is that it will not initialize CUDA
class NvmlCudaPlatform(CudaPlatformBase):

    @classmethod
    @lru_cache(maxsize=8)
    @with_nvml_context
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
        try:
            physical_device_id = device_id_to_physical_device_id(device_id)
            handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
            major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
            return DeviceCapability(major=major, minor=minor)
        except RuntimeError:
            return None

    @classmethod
    @lru_cache(maxsize=8)
    @with_nvml_context
    def has_device_capability(
        cls,
        capability: tuple[int, int] | int,
        device_id: int = 0,
    ) -> bool:
        try:
            return bool(super().has_device_capability(capability, device_id))
        except RuntimeError:
            return False

    @classmethod
    @lru_cache(maxsize=8)
    @with_nvml_context
    def get_device_name(cls, device_id: int = 0) -> str:
        physical_device_id = device_id_to_physical_device_id(device_id)
        return cls._get_physical_device_name(physical_device_id)

    @classmethod
    @lru_cache(maxsize=8)
    @with_nvml_context
    def get_device_uuid(cls, device_id: int = 0) -> str:
        physical_device_id = device_id_to_physical_device_id(device_id)
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return str(pynvml.nvmlDeviceGetUUID(handle))

    @classmethod
    @lru_cache(maxsize=8)
    @with_nvml_context
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        physical_device_id = device_id_to_physical_device_id(device_id)
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)

    @classmethod
    @with_nvml_context
    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
        """
        query if the set of gpus are fully connected by nvlink (1 hop)
        """
        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
                            handle,
                            peer_handle,
                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
                        )
                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                            return False
                    except pynvml.NVMLError:
                        logger.exception("NVLink detection failed. This is normal if"
                                         " your machine has no NVLink equipped.")
                        return False
        return True

    @classmethod
    def _get_physical_device_name(cls, device_id: int = 0) -> str:
        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
        return str(pynvml.nvmlDeviceGetName(handle))

    @classmethod
    @with_nvml_context
    def log_warnings(cls) -> None:
        device_ids: int = pynvml.nvmlDeviceGetCount()
        if device_ids > 1:
            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
            if (len(set(device_names)) > 1 and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"):
                logger.warning(
                    "Detected different devices in the system: %s. Please"
                    " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
                    "avoid unexpected behavior.",
                    ", ".join(device_names),
                )


class NonNvmlCudaPlatform(CudaPlatformBase):

    @classmethod
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)

    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        return str(torch.cuda.get_device_name(device_id))

    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        device_props = torch.cuda.get_device_properties(device_id)
        return int(device_props.total_memory)

    @classmethod
    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
        logger.exception("NVLink detection not possible, as context support was"
                         " not found. Assuming no NVLink available.")
        return False


# Autodetect either NVML-enabled or non-NVML platform
# based on whether NVML is available.
nvml_available = False
try:
    try:
        pynvml.nvmlInit()
        nvml_available = True
    except Exception:
        # On Jetson, NVML is not supported.
        nvml_available = False
finally:
    if nvml_available:
        pynvml.nvmlShutdown()

CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform

try:
    from sphinx.ext.autodoc.mock import _MockModule

    if not isinstance(pynvml, _MockModule):
        CudaPlatform.log_warnings()
except ModuleNotFoundError:
    CudaPlatform.log_warnings()