"""Utilities that depend on onnxruntime (ORT)."""

import functools
import logging
from pathlib import Path

from highlighter.core.cuda import is_nvidia_smi_available

LOGGER = logging.getLogger(__name__)


@functools.lru_cache(maxsize=1)
def is_cuda_provider_available() -> bool:
    """Return True if onnxruntime exposes CUDAExecutionProvider."""
    try:
        import onnxruntime as ort
    except Exception:
        return False

    try:
        providers = ort.get_available_providers()
    except Exception:
        return False
    return "CUDAExecutionProvider" in providers


def is_cuda_available() -> bool:
    """Return True when ORT has CUDA provider; logs if only NVIDIA tooling is found."""
    if is_cuda_provider_available():
        return True
    if is_nvidia_smi_available():
        LOGGER.debug("NVIDIA tooling detected but CUDAExecutionProvider is unavailable in onnxruntime.")
    return False


def download_model(
    url: str,
    path: str | Path,
    *,
    timeout: int = 30,
    check_cached: bool = False,
    show_progress: bool = False,
) -> Path:
    """Download a model artifact from URL with status checks and temp-file swap."""
    dest = Path(path)
    if check_cached and dest.exists():
        return dest

    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = dest.with_suffix(dest.suffix + ".partial")
    try:
        import requests
    except ImportError as exc:
        raise ImportError("requests is required to download ONNX models") from exc

    if show_progress:
        LOGGER.info("Downloading model from %s to %s", url, dest)
    else:
        LOGGER.debug("Downloading model from %s to %s", url, dest)

    try:
        with requests.get(url, stream=True, timeout=timeout) as response:
            response.raise_for_status()
            total_length = int(response.headers.get("content-length", 0))
            chunk_iter = response.iter_content(chunk_size=1024 * 1024)

            if show_progress:
                try:
                    import tqdm  # type: ignore

                    total = None
                    if total_length:
                        total = max(1, total_length // (1024 * 1024))
                    chunk_iter = tqdm.tqdm(
                        chunk_iter,
                        total=total,
                        bar_format="{l_bar}{bar:10}",
                    )
                except Exception as exc:
                    LOGGER.debug("tqdm unavailable; proceeding without progress bar: %s", exc)

            with tmp_path.open("wb") as f:
                for chunk in chunk_iter:
                    if chunk:
                        f.write(chunk)
            tmp_path.replace(dest)
    except Exception:
        if tmp_path.exists():
            tmp_path.unlink()
        raise

    return dest
