"""Pydantic models for hardware profiling."""

from pydantic import BaseModel, Field, computed_field


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

    name: str
    vram_gb: float
    vendor: str = "nvidia"  # nvidia, amd, intel
    cuda_version: str | None = None
    compute_capability: str | None = None
    driver_version: str | None = None


class HardwareProfile(BaseModel):
    """Complete hardware profile of the system."""

    cpu_cores: int
    cpu_name: str
    ram_total_gb: float
    ram_available_gb: float
    disk_free_gb: float
    gpus: list[GPUInfo] = Field(default_factory=list)
    os: str
    python_version: str
    torch_version: str | None = None
    cuda_available: bool = False

    @computed_field  # type: ignore[prop-decorator]
    @property
    def num_gpus(self) -> int:
        """Number of detected GPUs."""
        return len(self.gpus)

    @computed_field  # type: ignore[prop-decorator]
    @property
    def total_vram_gb(self) -> float:
        """Total VRAM across all GPUs."""
        return sum(gpu.vram_gb for gpu in self.gpus)

    @computed_field  # type: ignore[prop-decorator]
    @property
    def min_vram_gb(self) -> float:
        """VRAM of the smallest GPU (bottleneck for data parallel)."""
        if not self.gpus:
            return 0.0
        return min(gpu.vram_gb for gpu in self.gpus)
