"""GPU detection and profiling. Supports NVIDIA, AMD (ROCm), and Intel (XPU)."""

from bilgeai.profiler.models import GPUInfo


def detect_gpus() -> list[GPUInfo]:
    """Detect available GPUs. Tries pynvml (NVIDIA), then torch for all vendors."""
    all_gpus: list[GPUInfo] = []

    # NVIDIA via pynvml (fast, no torch needed)
    nvidia_gpus = _detect_via_pynvml()
    if nvidia_gpus is not None:
        all_gpus.extend(nvidia_gpus)
    else:
        # Fallback: NVIDIA via torch.cuda
        nvidia_torch = _detect_nvidia_via_torch()
        if nvidia_torch:
            all_gpus.extend(nvidia_torch)

    # AMD via torch ROCm
    amd_gpus = _detect_amd_via_torch()
    if amd_gpus:
        all_gpus.extend(amd_gpus)

    # Intel via torch XPU
    intel_gpus = _detect_intel_via_torch()
    if intel_gpus:
        all_gpus.extend(intel_gpus)

    # Apple Silicon via torch MPS
    mps_gpus = _detect_apple_mps()
    if mps_gpus:
        all_gpus.extend(mps_gpus)

    return all_gpus


def _detect_via_pynvml() -> list[GPUInfo] | None:
    """Detect NVIDIA GPUs using pynvml (no torch dependency)."""
    try:
        import pynvml

        pynvml.nvmlInit()
        device_count = pynvml.nvmlDeviceGetCount()
        gpus = []

        for i in range(device_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            name = pynvml.nvmlDeviceGetName(handle)
            if isinstance(name, bytes):
                name = name.decode("utf-8")

            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            vram_gb = round(mem_info.total / (1024**3), 1)

            try:
                driver = pynvml.nvmlSystemGetDriverVersion()
                if isinstance(driver, bytes):
                    driver = driver.decode("utf-8")
            except Exception:
                driver = None

            try:
                major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
                compute_cap = f"{major}.{minor}"
            except Exception:
                compute_cap = None

            gpus.append(
                GPUInfo(
                    name=name,
                    vram_gb=vram_gb,
                    cuda_version=_get_cuda_version_from_driver(),
                    compute_capability=compute_cap,
                    driver_version=driver,
                )
            )

        pynvml.nvmlShutdown()
        return gpus
    except (ImportError, Exception):
        return None


def _detect_nvidia_via_torch() -> list[GPUInfo] | None:
    """Detect NVIDIA GPUs using torch.cuda (fallback)."""
    try:
        import torch

        if not torch.cuda.is_available():
            return None

        gpus = []
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            gpus.append(
                GPUInfo(
                    name=props.name,
                    vram_gb=round(props.total_mem / (1024**3), 1),
                    vendor="nvidia",
                    cuda_version=torch.version.cuda,
                    compute_capability=f"{props.major}.{props.minor}",
                )
            )
        return gpus
    except (ImportError, Exception):
        return None


def _detect_amd_via_torch() -> list[GPUInfo] | None:
    """Detect AMD GPUs using torch with ROCm backend."""
    try:
        import torch

        if not hasattr(torch.version, "hip") or not torch.version.hip:
            return None
        if not torch.cuda.is_available():
            return None

        gpus = []
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            gpus.append(
                GPUInfo(
                    name=props.name,
                    vram_gb=round(props.total_mem / (1024**3), 1),
                    vendor="amd",
                )
            )
        return gpus
    except (ImportError, Exception):
        return None


def _detect_intel_via_torch() -> list[GPUInfo] | None:
    """Detect Intel GPUs using torch.xpu (Intel Extension for PyTorch)."""
    try:
        import torch

        if not hasattr(torch, "xpu") or not torch.xpu.is_available():
            return None

        gpus = []
        for i in range(torch.xpu.device_count()):
            props = torch.xpu.get_device_properties(i)
            name = props.name if hasattr(props, "name") else f"Intel XPU {i}"
            vram = (
                round(props.total_memory / (1024**3), 1) if hasattr(props, "total_memory") else 0.0
            )
            gpus.append(
                GPUInfo(
                    name=name,
                    vram_gb=vram,
                    vendor="intel",
                )
            )
        return gpus
    except (ImportError, Exception):
        return None


def _detect_apple_mps() -> list[GPUInfo] | None:
    """Detect Apple Silicon GPU using torch.backends.mps."""
    try:
        import torch

        if not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available():
            return None

        # Apple Silicon shares unified memory with CPU
        import psutil

        total_ram_gb = round(psutil.virtual_memory().total / (1024**3), 1)
        # MPS can use ~75% of unified memory for GPU tasks
        estimated_vram = round(total_ram_gb * 0.75, 1)

        # Detect chip name
        chip_name = _get_apple_chip_name()

        return [
            GPUInfo(
                name=chip_name or "Apple Silicon GPU",
                vram_gb=estimated_vram,
                vendor="apple",
            )
        ]
    except (ImportError, Exception):
        return None


def _get_apple_chip_name() -> str | None:
    """Try to detect Apple chip name (M1/M2/M3/M4)."""
    try:
        import platform

        if platform.system() != "Darwin":
            return None
        import subprocess

        result = subprocess.run(
            ["sysctl", "-n", "machdep.cpu.brand_string"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout.strip()
    except Exception:
        pass
    return None


def _get_cuda_version_from_driver() -> str | None:
    """Try to get CUDA version from nvml."""
    try:
        import pynvml

        cuda_version = pynvml.nvmlSystemGetCudaDriverVersion_v2()
        major = cuda_version // 1000
        minor = (cuda_version % 1000) // 10
        return f"{major}.{minor}"
    except Exception:
        return None
