"""Device resolution for acaua adapters.

Priority: `cuda > mps > cpu`. Explicit device strings are validated before
return so users get a clear error instead of a cryptic torch trace when they
request a device that isn't available.
"""

from __future__ import annotations

import torch


def resolve(device: str = "auto") -> torch.device:
    """Resolve a device string to a `torch.device`.

    Args:
        device: "auto" (default) picks the best available; explicit values
            ("cuda", "mps", "cpu", "cuda:1", ...) are validated and returned
            as-is.

    Raises:
        RuntimeError: explicit cuda/mps requested on a host that doesn't
            have the backend available.
    """
    if device == "auto":
        if torch.cuda.is_available():
            return torch.device("cuda")
        if _mps_available():
            return torch.device("mps")
        return torch.device("cpu")

    resolved = torch.device(device)
    if resolved.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError(
            f"device={device!r} requested but CUDA is not available on this host. "
            "Use device='auto' or device='cpu'."
        )
    if resolved.type == "mps" and not _mps_available():
        raise RuntimeError(f"device={device!r} requested but MPS is not available on this host.")
    return resolved


def _mps_available() -> bool:
    return hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
