"""Accelerator detection. ML libraries are imported only when installed."""

from __future__ import annotations

import functools
import importlib.util
from typing import Any


def _installed(name: str) -> bool:
    try:
        return importlib.util.find_spec(name) is not None
    except (ImportError, ValueError):
        return False


@functools.cache
def torch_gpu_available() -> bool:
    if not _installed("torch"):
        return False
    import torch

    return bool(torch.cuda.is_available())


def effective_accelerator(settings: Any) -> str:
    """What the image was built for, downgraded to cpu when no GPU is visible.

    ROCm builds of torch expose the GPU through the torch.cuda API.
    """
    if settings.accelerator in ("cuda", "rocm") and torch_gpu_available():
        return settings.accelerator
    return "cpu"


def torch_device(settings: Any) -> str:
    return "cuda" if effective_accelerator(settings) != "cpu" else "cpu"


def ctranslate2_device(settings: Any) -> str:
    # CTranslate2 ships CUDA builds only; ROCm hosts transcribe on the CPU.
    if settings.accelerator != "cuda":
        return "cpu"
    import ctranslate2

    return "cuda" if ctranslate2.get_cuda_device_count() > 0 else "cpu"


def onnx_providers(settings: Any) -> list[str]:
    import onnxruntime

    available = set(onnxruntime.get_available_providers())
    providers: list[str] = []
    if settings.accelerator == "cuda" and "CUDAExecutionProvider" in available:
        providers.append("CUDAExecutionProvider")
    if settings.accelerator == "rocm":
        for name in ("MIGraphXExecutionProvider", "ROCMExecutionProvider"):
            if name in available:
                providers.append(name)
    providers.append("CPUExecutionProvider")
    return providers