import torch


def resolve_torch_device(
    logger,
    requested_device: str = "auto",
    cuda_probe_fn=None,
):
    """Pick a safe torch device, falling back to CPU if CUDA is unavailable or unusable."""
    requested = (requested_device or "auto").strip().lower()
    valid = {"auto", "cpu", "cuda"}
    if requested not in valid:
        logger.warn(
            f"Unknown inference_device='{requested_device}', falling back to auto"
        )
        requested = "auto"

    if requested == "cpu":
        return "cpu"

    if not torch.cuda.is_available():
        if requested == "cuda":
            logger.warn("CUDA was requested but is not available; falling back to CPU")
        return "cpu"

    probe = cuda_probe_fn or _default_cuda_probe
    try:
        probe()
        return "cuda"
    except Exception as exc:
        logger.warn(
            "CUDA is visible but not usable by the installed PyTorch build; "
            f"falling back to CPU. Details: {exc}"
        )
        return "cpu"


def _default_cuda_probe():
    # Force a tiny kernel launch so unsupported GPU architectures fail early.
    tensor = torch.zeros(1, device="cuda")
    _ = tensor + 1
    torch.cuda.synchronize()
