# SPDX-License-Identifier: Apache-2.0

import gc
import os

import torch

import areal.utils.logging as logging

from .platform import Platform

logger = logging.getLogger("CUDAPlatform")


class CudaPlatform(Platform):
    device_name: str = "NVIDIA"
    device_type: str = "cuda"
    dispatch_key: str = "CUDA"
    ray_device_key: str = "GPU"
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
    ray_experimental_noset: str = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
    communication_backend: str = "nccl"

    def clear_memory(self) -> None:
        torch.cuda.synchronize()
        gc.collect()
        torch.cuda.empty_cache()

    @classmethod
    def clear_cublas_workspaces(cls) -> None:
        torch._C._cuda_clearCublasWorkspaces()

    @classmethod
    def get_vllm_worker_class(clas):
        try:
            from vllm.v1.worker.gpu_worker import Worker

            logger.info("Successfully imported vLLM V1 Worker.")
            return Worker
        except ImportError:
            pass

        try:
            from vllm.worker.worker import Worker

            logger.info("Successfully imported vLLM V0 Worker.")
            return Worker
        except ImportError as e:
            logger.error(
                "Failed to import vLLM Worker. "
                "Make sure vLLM is installed correctly: %s",
                e,
            )
            raise RuntimeError(
                "vLLM is not installed or not properly configured."
            ) from e

    @classmethod
    def set_allocator_settings(cls) -> None:
        torch.cuda.memory._set_allocator_settings("expandable_segments:False")

    @classmethod
    def set_numa_affinity(cls, local_rank: int) -> None:
        """Bind the current process to CPU cores local to the assigned GPU."""

        nvml_initialized = False
        try:
            import pynvml

            pynvml.nvmlInit()
            nvml_initialized = True
            handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank)
            pynvml.nvmlDeviceSetCpuAffinity(handle)
            cpu_set = os.sched_getaffinity(0)
            logger.info(
                "Set NUMA affinity for GPU %s: bound to %s CPU cores.",
                local_rank,
                len(cpu_set),
            )
        except ImportError:
            logger.warning(
                "pynvml (nvidia-ml-py) not available, skipping NUMA affinity setup."
            )
        except Exception as e:
            logger.warning("Failed to set NUMA affinity for GPU %s: %s", local_rank, e)
        finally:
            if nvml_initialized:
                pynvml.nvmlShutdown()

    @classmethod
    def get_custom_env_vars(cls) -> dict:
        env_vars = {
            # "RAY_DEBUG": "legacy"
            "TORCHINDUCTOR_COMPILE_THREADS": "2",
            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
        }
        return env_vars

    @classmethod
    def synchronize(cls) -> None:
        torch.cuda.synchronize()
