"""GPU type specifications."""

from __future__ import annotations

from dataclasses import dataclass


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

    ``price_per_hour`` is an approximate guide — actual billing uses
    live prices fetched from the Velar API at deployment time.
    """

    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).
# Prices reflect ~70 % markup over RunPod secure-cloud on-demand rates
# and may drift as RunPod adjusts pricing.  The backend always bills at
# the live rate returned by GET /api/v1/gpu/prices.

L4 = GPU(name="L4-24GB", vram_gb=24, price_per_hour=0.66)
RTX4090 = GPU(name="RTX4090-24GB", vram_gb=24, price_per_hour=1.00)
L40S = GPU(name="L40S-48GB", vram_gb=48, price_per_hour=1.46)
A100 = GPU(name="A100-80GB", vram_gb=80, price_per_hour=2.36)
H100 = GPU(name="H100-80GB", vram_gb=80, price_per_hour=4.06)
H100_SXM = GPU(name="H100-SXM-80GB", vram_gb=80, price_per_hour=4.57)
H200 = GPU(name="H200-141GB", vram_gb=141, price_per_hour=6.10)

GPU_MAP: dict[str, GPU] = {
    # L4
    "L4": L4,
    "L4-24GB": L4,
    # RTX 4090
    "RTX4090": RTX4090,
    "RTX4090-24GB": RTX4090,
    "RTX 4090": RTX4090,
    # L40S
    "L40S": L40S,
    "L40S-48GB": L40S,
    # A100
    "A100": A100,
    "A100-80GB": A100,
    # H100 PCIe
    "H100": H100,
    "H100-80GB": H100,
    # H100 SXM
    "H100-SXM": H100_SXM,
    "H100-SXM-80GB": H100_SXM,
    "H100 SXM": H100_SXM,
    # H200
    "H200": H200,
    "H200-141GB": H200,
}


def resolve(spec: str | GPU) -> GPU:
    """Resolve a GPU specification to a GPU object."""
    if isinstance(spec, GPU):
        return spec
    key = spec.strip()
    gpu = GPU_MAP.get(key) or GPU_MAP.get(key.upper())
    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
