"""Device registry loader — reads devices.toml and returns DeviceInfo objects."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import toml

from tests.remote.runner import DeviceInfo

REGISTRY_PATH = Path(__file__).parent / "devices.toml"


@dataclass
class DeviceSpec:
    """Full device specification from the registry (extends DeviceInfo with hardware info)."""

    name: str
    hostname: str
    username: str
    shell: str
    os: str
    gpu: str
    ram_gb: int
    ssh_key: str | None = None
    port: int = 22
    description: str = ""

    def to_device_info(self) -> DeviceInfo:
        """Convert to the minimal DeviceInfo needed by RemoteTestRunner."""
        return DeviceInfo(
            hostname=self.hostname,
            username=self.username,
            shell=self.shell,
            os=self.os,
            ssh_key=self.ssh_key,
            port=self.port,
        )


def load_registry(path: Path | None = None) -> dict[str, DeviceSpec]:
    """Load all devices from the TOML registry.

    Args:
        path: Path to the TOML file. Defaults to devices.toml in this directory.

    Returns:
        Dict mapping device name to DeviceSpec.
    """
    registry_path = path or REGISTRY_PATH

    if not registry_path.exists():
        raise FileNotFoundError(f"Device registry not found: {registry_path}")

    data = toml.load(registry_path)
    devices: dict[str, DeviceSpec] = {}

    for name, config in data.items():
        devices[name] = DeviceSpec(
            name=name,
            hostname=config["hostname"],
            username=config["username"],
            shell=config["shell"],
            os=config["os"],
            gpu=config.get("gpu", "unknown"),
            ram_gb=config.get("ram_gb", 0),
            ssh_key=config.get("ssh_key"),
            port=config.get("port", 22),
            description=config.get("description", ""),
        )

    return devices


def get_device(name: str, path: Path | None = None) -> DeviceSpec:
    """Load a single device by name from the registry.

    Raises:
        KeyError: If the device name is not found.
    """
    registry = load_registry(path)
    if name not in registry:
        available = ", ".join(sorted(registry.keys()))
        raise KeyError(f"Device '{name}' not found. Available: {available}")
    return registry[name]
