"""GPU detection utilities.

This module provides utilities for detecting available GPUs
on the system from various vendors (NVIDIA, AMD, Intel).
"""

from __future__ import annotations

import re
import shutil
import subprocess
from pathlib import Path

from gaol.gpu.models import GPUInfo, GPUSupport, GPUVendor


class GPUDetector:
    """Detect available GPUs on the system.

    Provides static methods for detecting GPUs from different vendors
    and checking runtime availability.

    Example:
        >>> gpus = GPUDetector.detect_all()
        >>> for gpu in gpus:
        ...     print(f"{gpu.vendor}: {gpu.name}")
        nvidia: NVIDIA GeForce RTX 3080
    """

    @staticmethod
    def detect_all() -> list[GPUInfo]:
        """Detect all GPUs from all vendors.

        Returns:
            List of detected GPUs from all vendors.
        """
        gpus: list[GPUInfo] = []
        gpus.extend(GPUDetector.detect_nvidia())
        gpus.extend(GPUDetector.detect_amd())
        gpus.extend(GPUDetector.detect_intel())
        return gpus

    @staticmethod
    def detect_nvidia() -> list[GPUInfo]:
        """Detect NVIDIA GPUs using nvidia-smi.

        Returns:
            List of detected NVIDIA GPUs.
        """
        if not shutil.which("nvidia-smi"):
            return []

        try:
            result = subprocess.run(
                [
                    "nvidia-smi",
                    "--query-gpu=index,name,uuid,memory.total,pci.bus_id",
                    "--format=csv,noheader,nounits",
                ],
                capture_output=True,
                text=True,
                timeout=10,
            )

            if result.returncode != 0:
                return []

            gpus: list[GPUInfo] = []
            for line in result.stdout.strip().split("\n"):
                if not line.strip():
                    continue

                parts = [p.strip() for p in line.split(",")]
                if len(parts) >= 5:
                    try:
                        memory_mb = int(float(parts[3])) if parts[3] else None
                    except ValueError:
                        memory_mb = None

                    gpus.append(
                        GPUInfo(
                            index=int(parts[0]),
                            name=parts[1],
                            vendor=GPUVendor.NVIDIA,
                            uuid=parts[2] if parts[2] else None,
                            memory_mb=memory_mb,
                            pci_bus_id=parts[4] if parts[4] else None,
                        )
                    )

            return gpus

        except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
            return []

    @staticmethod
    def detect_amd() -> list[GPUInfo]:
        """Detect AMD GPUs using rocm-smi or lspci.

        Returns:
            List of detected AMD GPUs.
        """
        gpus: list[GPUInfo] = []

        # Try rocm-smi first
        if shutil.which("rocm-smi"):
            try:
                result = subprocess.run(
                    ["rocm-smi", "--showproductname"],
                    capture_output=True,
                    text=True,
                    timeout=10,
                )

                if result.returncode == 0:
                    # Parse rocm-smi output
                    lines = result.stdout.strip().split("\n")
                    index = 0
                    for line in lines:
                        # Look for GPU lines (usually "GPU[X]" pattern)
                        match = re.search(r"GPU\[(\d+)\].*?:\s*(.+)", line)
                        if match:
                            gpus.append(
                                GPUInfo(
                                    index=int(match.group(1)),
                                    name=match.group(2).strip(),
                                    vendor=GPUVendor.AMD,
                                )
                            )
                            index += 1

                    if gpus:
                        return gpus
            except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
                pass

        # Fall back to checking /dev/dri for AMD render nodes
        # and parsing lspci
        if shutil.which("lspci"):
            try:
                result = subprocess.run(
                    ["lspci", "-nn"],
                    capture_output=True,
                    text=True,
                    timeout=10,
                )

                if result.returncode == 0:
                    index = 0
                    for line in result.stdout.split("\n"):
                        # AMD/ATI VGA or 3D controller
                        if "AMD" in line or "ATI" in line:
                            if "VGA" in line or "3D" in line or "Display" in line:
                                # Extract PCI address and name
                                pci_match = re.match(r"(\S+)\s+", line)
                                name_match = re.search(
                                    r":\s*(.+?)\s*\[", line
                                ) or re.search(r":\s*(.+)", line)

                                gpus.append(
                                    GPUInfo(
                                        index=index,
                                        name=name_match.group(1).strip()
                                        if name_match
                                        else "AMD GPU",
                                        vendor=GPUVendor.AMD,
                                        pci_bus_id=f"0000:{pci_match.group(1)}"
                                        if pci_match
                                        else None,
                                    )
                                )
                                index += 1
            except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
                pass

        return gpus

    @staticmethod
    def detect_intel() -> list[GPUInfo]:
        """Detect Intel GPUs from /dev/dri and lspci.

        Returns:
            List of detected Intel GPUs.
        """
        gpus: list[GPUInfo] = []

        # Check for Intel render devices
        dri_path = Path("/dev/dri")
        if not dri_path.exists():
            return []

        # Use lspci to identify Intel GPUs
        if shutil.which("lspci"):
            try:
                result = subprocess.run(
                    ["lspci", "-nn"],
                    capture_output=True,
                    text=True,
                    timeout=10,
                )

                if result.returncode == 0:
                    index = 0
                    for line in result.stdout.split("\n"):
                        if "Intel" in line and (
                            "VGA" in line or "3D" in line or "Display" in line
                        ):
                            pci_match = re.match(r"(\S+)\s+", line)
                            name_match = re.search(
                                r":\s*(.+?)\s*\[", line
                            ) or re.search(r":\s*(.+)", line)

                            gpus.append(
                                GPUInfo(
                                    index=index,
                                    name=name_match.group(1).strip()
                                    if name_match
                                    else "Intel GPU",
                                    vendor=GPUVendor.INTEL,
                                    pci_bus_id=f"0000:{pci_match.group(1)}"
                                    if pci_match
                                    else None,
                                )
                            )
                            index += 1
            except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
                pass

        return gpus

    @staticmethod
    def get_nvidia_driver_version() -> str | None:
        """Get the installed NVIDIA driver version.

        Returns:
            Driver version string or None if not available.
        """
        if not shutil.which("nvidia-smi"):
            return None

        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
                capture_output=True,
                text=True,
                timeout=10,
            )

            if result.returncode == 0:
                version = result.stdout.strip().split("\n")[0].strip()
                return version if version else None
        except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
            pass

        return None

    @staticmethod
    def get_cuda_version() -> str | None:
        """Get the installed CUDA version.

        Returns:
            CUDA version string or None if not available.
        """
        if not shutil.which("nvidia-smi"):
            return None

        try:
            result = subprocess.run(
                ["nvidia-smi"],
                capture_output=True,
                text=True,
                timeout=10,
            )

            if result.returncode == 0:
                # Look for "CUDA Version: X.Y" in output
                match = re.search(r"CUDA Version:\s*(\S+)", result.stdout)
                if match:
                    return match.group(1)
        except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
            pass

        return None

    @staticmethod
    def is_nvidia_runtime_available() -> bool:
        """Check if nvidia-container-runtime is installed.

        Returns:
            True if nvidia-container-runtime is available.
        """
        return shutil.which("nvidia-container-runtime") is not None

    @staticmethod
    def is_rocm_available() -> bool:
        """Check if ROCm is installed.

        Returns:
            True if ROCm tools are available.
        """
        return shutil.which("rocm-smi") is not None

    @staticmethod
    def get_rocm_version() -> str | None:
        """Get the installed ROCm version.

        Returns:
            ROCm version string or None if not available.
        """
        if not shutil.which("rocm-smi"):
            return None

        try:
            result = subprocess.run(
                ["rocm-smi", "--showversion"],
                capture_output=True,
                text=True,
                timeout=10,
            )

            if result.returncode == 0:
                # Look for version in output
                match = re.search(r"version\s*:?\s*(\S+)", result.stdout, re.IGNORECASE)
                if match:
                    return match.group(1)
        except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
            pass

        return None

    @staticmethod
    def get_support_info() -> GPUSupport:
        """Get comprehensive GPU support information.

        Returns:
            GPUSupport object with all detected GPUs and capabilities.
        """
        gpus = GPUDetector.detect_all()

        nvidia_gpus = [g for g in gpus if g.vendor == GPUVendor.NVIDIA]
        amd_gpus = [g for g in gpus if g.vendor == GPUVendor.AMD]
        intel_gpus = [g for g in gpus if g.vendor == GPUVendor.INTEL]

        return GPUSupport(
            nvidia_available=len(nvidia_gpus) > 0,
            nvidia_driver_version=GPUDetector.get_nvidia_driver_version(),
            nvidia_runtime_available=GPUDetector.is_nvidia_runtime_available(),
            cuda_version=GPUDetector.get_cuda_version(),
            amd_available=len(amd_gpus) > 0,
            rocm_available=GPUDetector.is_rocm_available(),
            rocm_version=GPUDetector.get_rocm_version(),
            intel_available=len(intel_gpus) > 0,
            gpus=gpus,
        )
