"""GPU type specifications."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class GPU:
    """GPU type specification."""

    name: str
    vram_gb: int
    price_per_hour: float

    def __str__(self) -> str:
        return self.name


# Available GPU types (must match backend models.GPUType constants)
A100 = GPU(name="A100-80GB", vram_gb=80, price_per_hour=3.20)
A10 = GPU(name="A10-24GB", vram_gb=24, price_per_hour=1.40)
L4 = GPU(name="L4-24GB", vram_gb=24, price_per_hour=0.85)

GPU_MAP: dict[str, GPU] = {
    "A100": A100,
    "A100-80GB": A100,
    "A10": A10,
    "A10-24GB": A10,
    "L4": L4,
    "L4-24GB": L4,
}


def resolve(spec: str | GPU) -> GPU:
    """Resolve a GPU specification to a GPU object."""
    if isinstance(spec, GPU):
        return spec
    gpu = GPU_MAP.get(spec.upper()) or GPU_MAP.get(spec)
    if gpu is None:
        available = ", ".join(sorted(set(g.name for g in GPU_MAP.values())))
        raise ValueError(f"Unknown GPU type: {spec}. Available: {available}")
    return gpu
