"""GPU configuration models.

This module provides Pydantic models for GPU configuration and
device mapping used in container GPU passthrough.
"""

from __future__ import annotations

from enum import Enum

from pydantic import BaseModel, Field


class GPUVendor(str, Enum):
    """GPU vendor types."""

    NVIDIA = "nvidia"
    AMD = "amd"
    INTEL = "intel"
    AUTO = "auto"  # Auto-detect


class GPUInfo(BaseModel):
    """Information about a detected GPU.

    Attributes:
        index: GPU device index (0, 1, 2, etc.)
        name: Human-readable GPU name
        vendor: GPU vendor (nvidia, amd, intel)
        uuid: Unique identifier for the GPU
        memory_mb: Total GPU memory in megabytes
        pci_bus_id: PCI bus address for passthrough
        driver_version: Installed driver version
        cuda_version: CUDA version (NVIDIA only)
    """

    index: int
    name: str
    vendor: GPUVendor
    uuid: str | None = None
    memory_mb: int | None = None
    pci_bus_id: str | None = None
    driver_version: str | None = None
    cuda_version: str | None = None


class GPUConfig(BaseModel):
    """GPU configuration for containers.

    Attributes:
        enabled: Whether GPU passthrough is enabled
        vendor: GPU vendor to use (auto-detect if not specified)
        device_ids: List of GPU device IDs to pass through (empty = all)
        capabilities: NVIDIA capabilities (compute, utility, graphics, etc.)
        require_nvidia_runtime: Use nvidia-container-runtime for Docker
        memory_limit: Memory limit per GPU (e.g., "8g")

    Example:
        >>> config = GPUConfig(enabled=True, device_ids=[0, 1])
        >>> config.vendor
        <GPUVendor.AUTO: 'auto'>
    """

    enabled: bool = False
    vendor: GPUVendor = GPUVendor.AUTO
    device_ids: list[int] = Field(default_factory=list)
    capabilities: list[str] = Field(default_factory=lambda: ["compute", "utility"])
    require_nvidia_runtime: bool = True
    memory_limit: str | None = None


class DeviceMapping(BaseModel):
    """Generic device mapping for passthrough.

    Attributes:
        host_path: Device path on host (e.g., /dev/nvidia0)
        container_path: Device path in container (defaults to host_path)
        permissions: Device permissions (r=read, w=write, m=mknod)

    Example:
        >>> device = DeviceMapping(host_path="/dev/nvidia0")
        >>> device.container_path  # Defaults to host_path
        '/dev/nvidia0'
    """

    host_path: str
    container_path: str | None = None
    permissions: str = "rwm"

    @property
    def effective_container_path(self) -> str:
        """Get the effective container path (defaults to host_path)."""
        return self.container_path or self.host_path


class GPUSupport(BaseModel):
    """Information about GPU support on the system.

    Attributes:
        nvidia_available: Whether NVIDIA GPUs are detected
        nvidia_driver_version: NVIDIA driver version
        nvidia_runtime_available: Whether nvidia-container-runtime is installed
        amd_available: Whether AMD GPUs are detected
        rocm_available: Whether ROCm is installed
        intel_available: Whether Intel GPUs are detected
        gpus: List of detected GPUs
    """

    nvidia_available: bool = False
    nvidia_driver_version: str | None = None
    nvidia_runtime_available: bool = False
    cuda_version: str | None = None
    amd_available: bool = False
    rocm_available: bool = False
    rocm_version: str | None = None
    intel_available: bool = False
    gpus: list[GPUInfo] = Field(default_factory=list)
