"""Shared device-resolution and model-placement helpers."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any


if TYPE_CHECKING:
    import torch


logger = logging.getLogger(__name__)


def resolve_device(preferred: str = "cuda") -> torch.device:
    """Select the best available compute device.

    Args:
        preferred: Requested device key (``cuda``, ``mps``, or ``cpu``).

    Returns:
        Resolved ``torch.device`` value.
    """
    import torch

    if preferred == "cuda" and torch.cuda.is_available():
        return torch.device("cuda")

    if preferred == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")

    if preferred in ("cuda", "mps"):
        logger.warning(
            "Requested device '%s' not available — falling back to CPU",
            preferred,
        )

    return torch.device("cpu")


def move_to_device(model: Any, device: torch.device) -> torch.device:
    """Move a model to the target device.

    Args:
        model: Torch model or wrapper object with ``.model`` attribute.
        device: Target device.

    Returns:
        The device used for placement.
    """
    if hasattr(model, "model"):
        model.model = model.model.to(device)
        if hasattr(model, "device"):
            model.device = device
    elif hasattr(model, "to"):
        model.to(device)

    return device


def resolve_torch_device() -> torch.device:
    """Resolve the best available torch inference device.

    Args:
        None.

    Returns:
        ``torch.device("cuda")`` when available, otherwise CPU.
    """
    import torch

    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def prepare_torch_model_for_inference(
    model: torch.nn.Module,
    device: torch.device | None = None,
) -> torch.nn.Module:
    """Move model to device and set eval mode for inference.

    Args:
        model: Torch module to prepare.
        device: Optional explicit target device.

    Returns:
        Inference-ready model on the resolved device.
    """
    target_device = device or resolve_torch_device()
    model = model.to(target_device)
    model.eval()
    return model
