engine/hardware.py

"""
Describing a machine in terms the maths cares about.

The whole job here is to turn "I have a Windows laptop with an RTX 3060 and
16 GB of RAM" into two numbers the advisor can reason about:

  - fast_budget_gb : memory the model can use *on the fast path* (the GPU, or
                     on Apple Silicon, the shared memory the GPU can borrow)
  - total_budget_gb: the absolute most a model can use if we let it spill onto
                     ordinary RAM (slower, but it runs)

Everything is deliberately conservative. We'd rather say "this might be tight"
and be wrong than promise something that then fails to load.
"""

from dataclasses import dataclass


# --------------------------------------------------------------------------
# Common consumer GPUs -> VRAM (GB). So people pick a name, not a number.
# VRAM is baked into the label too, because some cards ship in two sizes.
# --------------------------------------------------------------------------
GPU_PRESETS: dict[str, float] = {
    # NVIDIA RTX 50-series
    "NVIDIA RTX 5090 (32 GB)": 32,
    "NVIDIA RTX 5080 (16 GB)": 16,
    "NVIDIA RTX 5070 Ti (16 GB)": 16,
    "NVIDIA RTX 5070 (12 GB)": 12,
    "NVIDIA RTX 5060 Ti (16 GB)": 16,
    "NVIDIA RTX 5060 (8 GB)": 8,
    # NVIDIA RTX 40-series
    "NVIDIA RTX 4090 (24 GB)": 24,
    "NVIDIA RTX 4080 (16 GB)": 16,
    "NVIDIA RTX 4070 Ti (12 GB)": 12,
    "NVIDIA RTX 4070 (12 GB)": 12,
    "NVIDIA RTX 4060 Ti (16 GB)": 16,
    "NVIDIA RTX 4060 (8 GB)": 8,
    # NVIDIA RTX 30-series
    "NVIDIA RTX 3090 (24 GB)": 24,
    "NVIDIA RTX 3080 (10 GB)": 10,
    "NVIDIA RTX 3070 (8 GB)": 8,
    "NVIDIA RTX 3060 (12 GB)": 12,
    "NVIDIA RTX 3050 (8 GB)": 8,
    # Older / budget NVIDIA
    "NVIDIA GTX 1660 (6 GB)": 6,
    "NVIDIA GTX 1650 (4 GB)": 4,
    # AMD
    "AMD RX 7900 XTX (24 GB)": 24,
    "AMD RX 7800 XT (16 GB)": 16,
    "AMD RX 7600 (8 GB)": 8,
    "AMD RX 6700 XT (12 GB)": 12,
    # Laptop integrated (no real VRAM — uses shared system RAM)
    "Intel built-in graphics (no separate card)": 0,
    "AMD built-in graphics (no separate card)": 0,
}

# Apple Silicon: there's no separate VRAM. The GPU shares system memory, and
# macOS lets it borrow a large slice. We treat it specially below.
APPLE_CHIPS: dict[str, int] = {
    "Apple M1 / M2 / M3 / M4 (base)": 8,    # default RAM if they don't know
    "Apple M-series Pro": 16,
    "Apple M-series Max": 32,
    "Apple M-series Ultra": 64,
}


@dataclass
class HardwareSpec:
    """A machine, described just enough to reason about it."""

    os: str = "windows"               # windows | macos | linux
    ram_gb: float = 16.0              # system RAM
    gpu_vendor: str = "none"          # nvidia | amd | apple | intel | none
    vram_gb: float = 0.0              # dedicated GPU memory (0 if shared/none)
    is_apple_silicon: bool = False
    gpu_label: str = "No dedicated graphics card"
    form_factor: str = "laptop"       # laptop | desktop | mac | sbc

    # -- derived memory budgets -------------------------------------------

    @property
    def fast_budget_gb(self) -> float:
        """Memory available on the *fast* path (GPU / Apple shared memory)."""
        if self.is_apple_silicon:
            # macOS lets the GPU use a large fraction of unified memory.
            # ~70% is a safe, widely-quoted working figure.
            return round(self.ram_gb * 0.70, 1)
        if self.gpu_vendor in ("nvidia", "amd") and self.vram_gb > 0:
            # Leave headroom for the display, driver, and other apps.
            return round(self.vram_gb * 0.85, 1)
        # Integrated graphics / CPU-only: no meaningful fast path.
        return 0.0

    @property
    def os_reserve_gb(self) -> float:
        """RAM we set aside for the operating system + other open programs.

        Windows idles heavy; a headless Raspberry Pi barely uses anything.
        Being honest here matters: too small a reserve over-promises.
        """
        return {
            "sbc": 1.0,
            "mac": 3.0,
            "desktop": 3.0,
            "laptop": 3.5,
        }.get(self.form_factor, 3.0) if self.os != "linux" else {
            "sbc": 1.0,
        }.get(self.form_factor, 2.0)

    @property
    def total_budget_gb(self) -> float:
        """The most a model can use if it spills onto ordinary RAM (slower)."""
        if self.is_apple_silicon:
            return self.fast_budget_gb  # unified memory — same pool
        # Dedicated VRAM (fully usable on the fast path) PLUS a conservative
        # slice of system RAM for CPU offload, after reserving room for the OS.
        ram_for_model = max(0.0, self.ram_gb - self.os_reserve_gb) * 0.9
        return round(self.vram_gb + ram_for_model, 1)

    @property
    def has_fast_path(self) -> bool:
        return self.fast_budget_gb >= 1.0


def build_spec(
    *,
    computer_kind: str,
    ram_gb: float,
    gpu_choice: str,
    apple_chip: str | None = None,
) -> HardwareSpec:
    """Turn friendly UI selections into a HardwareSpec.

    computer_kind: "Windows laptop/desktop", "Mac", "Linux PC",
                   "Raspberry Pi / mini PC"
    gpu_choice:    a key from GPU_PRESETS, or one of the "don't know" options.
    """
    kind = computer_kind.lower()

    # ---- Mac / Apple Silicon -------------------------------------------
    if "mac" in kind:
        chip = apple_chip or "Apple M1 / M2 / M3 / M4 (base)"
        return HardwareSpec(
            os="macos",
            ram_gb=ram_gb,
            gpu_vendor="apple",
            vram_gb=0.0,
            is_apple_silicon=True,
            gpu_label=f"{chip} (shares your {ram_gb:g} GB of memory)",
            form_factor="mac",
        )

    # ---- Raspberry Pi / tiny single-board ------------------------------
    if "raspberry" in kind or "mini" in kind or "sbc" in kind:
        return HardwareSpec(
            os="linux",
            ram_gb=ram_gb,
            gpu_vendor="none",
            vram_gb=0.0,
            gpu_label="No dedicated graphics card (tiny computer)",
            form_factor="sbc",
        )

    # ---- Windows / Linux PC with a possible discrete GPU ---------------
    os_name = "linux" if "linux" in kind else "windows"
    form = "desktop" if "desktop" in kind else "laptop"

    vram = GPU_PRESETS.get(gpu_choice, 0.0)
    if "nvidia" in gpu_choice.lower():
        vendor = "nvidia"
    elif "amd" in gpu_choice.lower() and "built-in" not in gpu_choice.lower():
        vendor = "amd"
    elif "built-in" in gpu_choice.lower():
        vendor = "intel" if "intel" in gpu_choice.lower() else "amd"
    else:
        vendor = "none"

    label = gpu_choice if vram > 0 else "No dedicated graphics card (built-in graphics only)"

    return HardwareSpec(
        os=os_name,
        ram_gb=ram_gb,
        gpu_vendor=vendor,
        vram_gb=vram,
        is_apple_silicon=False,
        gpu_label=label,
        form_factor=form,
    )