"""NVIDIA GPU configuration handler.

This module provides NVIDIA-specific GPU configuration for different
container runtimes (Docker, nspawn, libvirt).
"""

from __future__ import annotations

from pathlib import Path

from gaol.gpu.models import DeviceMapping, GPUConfig


class NvidiaGPU:
    """NVIDIA GPU configuration handler.

    Provides methods to generate runtime-specific GPU configurations
    for NVIDIA GPUs using nvidia-container-toolkit.

    Example:
        >>> config = GPUConfig(enabled=True, device_ids=[0])
        >>> docker_config = NvidiaGPU.get_docker_config(config)
        >>> docker_config["device_requests"][0]["driver"]
        'nvidia'
    """

    # Standard NVIDIA device paths
    CONTROL_DEVICES = [
        "/dev/nvidiactl",
        "/dev/nvidia-uvm",
        "/dev/nvidia-uvm-tools",
        "/dev/nvidia-modeset",
    ]

    @staticmethod
    def get_docker_config(gpu_config: GPUConfig) -> dict:
        """Get Docker-specific GPU configuration.

        Uses Docker's device_requests API which works with
        nvidia-container-toolkit (Docker >= 19.03).

        Args:
            gpu_config: GPU configuration.

        Returns:
            Dictionary with device_requests and environment variables.
        """
        # Determine device IDs
        if gpu_config.device_ids:
            device_ids = [str(d) for d in gpu_config.device_ids]
        else:
            device_ids = ["all"]

        # Build device request
        device_requests = [
            {
                "Driver": "nvidia",
                "DeviceIDs": device_ids,
                "Capabilities": [gpu_config.capabilities],
            }
        ]

        # Environment variables
        environment = {
            "NVIDIA_VISIBLE_DEVICES": ",".join(device_ids)
            if device_ids != ["all"]
            else "all",
            "NVIDIA_DRIVER_CAPABILITIES": ",".join(gpu_config.capabilities),
        }

        # Add memory limit if specified
        if gpu_config.memory_limit:
            environment["NVIDIA_MEM_PER_GPU"] = gpu_config.memory_limit

        return {
            "device_requests": device_requests,
            "environment": environment,
        }

    @staticmethod
    def get_nspawn_config(gpu_config: GPUConfig) -> dict:
        """Get nspawn-specific GPU configuration.

        For systemd-nspawn, we need to bind mount the NVIDIA device files
        and set environment variables.

        Args:
            gpu_config: GPU configuration.

        Returns:
            Dictionary with bind mounts and environment variables.
        """
        binds: list[str] = []

        # Add per-GPU device files
        if gpu_config.device_ids:
            for device_id in gpu_config.device_ids:
                device_path = f"/dev/nvidia{device_id}"
                if Path(device_path).exists():
                    binds.append(device_path)
        else:
            # Add all nvidia devices
            for i in range(16):  # Support up to 16 GPUs
                device_path = f"/dev/nvidia{i}"
                if Path(device_path).exists():
                    binds.append(device_path)

        # Add control devices
        for device in NvidiaGPU.CONTROL_DEVICES:
            if Path(device).exists():
                binds.append(device)

        # Environment variables
        if gpu_config.device_ids:
            visible_devices = ",".join(str(d) for d in gpu_config.device_ids)
        else:
            visible_devices = "all"

        environment = {
            "NVIDIA_VISIBLE_DEVICES": visible_devices,
            "NVIDIA_DRIVER_CAPABILITIES": ",".join(gpu_config.capabilities),
        }

        return {
            "binds": binds,
            "environment": environment,
        }

    @staticmethod
    def get_libvirt_config(
        gpu_config: GPUConfig,  # noqa: ARG004
        pci_addresses: list[str],
    ) -> dict:
        """Get libvirt-specific GPU configuration.

        For libvirt, we generate PCI passthrough XML for each GPU.

        Args:
            gpu_config: GPU configuration (reserved for future use).
            pci_addresses: List of PCI addresses (e.g., "0000:05:00.0").

        Returns:
            Dictionary with XML snippets and configuration.
        """
        xml_parts: list[str] = []

        for pci_addr in pci_addresses:
            # Parse PCI address: domain:bus:slot.function
            parts = pci_addr.replace(":", ".").split(".")
            if len(parts) >= 4:
                domain = parts[0]
                bus = parts[1]
                slot = parts[2]
                function = parts[3]

                xml_parts.append(
                    f"""    <hostdev mode='subsystem' type='pci' managed='yes'>
      <source>
        <address domain='0x{domain}' bus='0x{bus}' slot='0x{slot}' function='0x{function}'/>
      </source>
    </hostdev>"""
                )

        return {
            "hostdev_xml": "\n".join(xml_parts),
            "requires_iommu": True,
        }

    @staticmethod
    def get_required_devices(device_ids: list[int] | None = None) -> list[str]:
        """Get list of required NVIDIA device paths.

        Args:
            device_ids: Specific GPU device IDs, or None for all.

        Returns:
            List of device paths.
        """
        devices: list[str] = []

        # Add per-GPU devices
        if device_ids:
            for device_id in device_ids:
                devices.append(f"/dev/nvidia{device_id}")
        else:
            # Add all that exist
            for i in range(16):
                device_path = f"/dev/nvidia{i}"
                if Path(device_path).exists():
                    devices.append(device_path)

        # Add control devices
        for device in NvidiaGPU.CONTROL_DEVICES:
            if Path(device).exists():
                devices.append(device)

        return devices

    @staticmethod
    def get_device_mappings(
        device_ids: list[int] | None = None,
    ) -> list[DeviceMapping]:
        """Get device mappings for NVIDIA GPUs.

        Args:
            device_ids: Specific GPU device IDs, or None for all.

        Returns:
            List of DeviceMapping objects.
        """
        mappings: list[DeviceMapping] = []

        for device_path in NvidiaGPU.get_required_devices(device_ids):
            mappings.append(
                DeviceMapping(
                    host_path=device_path,
                    container_path=device_path,
                    permissions="rwm",
                )
            )

        return mappings

    @staticmethod
    def get_required_capabilities() -> list[str]:
        """Get Linux capabilities required for GPU access.

        Returns:
            List of capability names.
        """
        return ["SYS_ADMIN"]
