"""I/O and runtime helpers for prediction workflows."""

from __future__ import annotations

import tempfile
from contextlib import contextmanager, suppress
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse

import requests
import torch

from matrice_models.common.device import prepare_torch_model_for_inference as common_prepare_torch_model_for_inference
from matrice_models.common.device import resolve_torch_device as common_resolve_torch_device


if TYPE_CHECKING:
    from collections.abc import Generator


def is_remote_path(path: str) -> bool:
    """Return whether a path points to an HTTP(S) resource.

    Args:
        path: Input path or URL string to inspect.

    Returns:
        ``True`` when the scheme is HTTP or HTTPS, otherwise ``False``.
    """
    parsed = urlparse(path)
    return parsed.scheme in {"http", "https"}


def download_remote_file(
    url: str,
    *,
    suffix: str = ".pt",
    timeout_seconds: int = 30,
    retries: int = 2,
) -> str:
    """Download a remote artifact into a temporary file and return its local path.

    The file is intentionally not auto-deleted because model runtimes can keep
    lazy references to checkpoint files after constructor return.

    Args:
        url: Remote URL for the artifact to download.
        suffix: Suffix used when creating temporary file name.
        timeout_seconds: Per-request timeout passed to ``requests.get``.
        retries: Number of retry attempts after the initial try.

    Returns:
        Local filesystem path to the downloaded temporary file.
    """
    last_error: Exception | None = None
    for _ in range(max(1, retries + 1)):
        tmp_path: str | None = None
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_handle:
                tmp_path = tmp_handle.name
                response = requests.get(url, timeout=timeout_seconds, stream=True)
                response.raise_for_status()
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        tmp_handle.write(chunk)
        except Exception as exc:  # pragma: no cover - retry path
            last_error = exc
            if tmp_path:
                with suppress(OSError):
                    Path(tmp_path).unlink()
        else:
            return tmp_path

    assert last_error is not None
    raise last_error


@contextmanager
def temporary_file_from_bytes(data: bytes, *, suffix: str = ".bin") -> Generator[str, None, None]:
    """Write bytes to a temporary file and clean it up on exit.

    Args:
        data: Raw bytes to write into the temporary file.
        suffix: Filename suffix used for the temporary file.

    Returns:
        Context manager yielding the temporary file path.
    """
    tmp_path: str | None = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as handle:
            handle.write(data)
            tmp_path = handle.name
        yield tmp_path
    finally:
        if tmp_path and Path(tmp_path).exists():
            with suppress(OSError):
                Path(tmp_path).unlink()


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

    Args:
        None.

    Returns:
        ``torch.device("cuda")`` when CUDA is available, else CPU device.
    """
    return common_resolve_torch_device()


def prepare_torch_model_for_inference(model: torch.nn.Module, device: torch.device | None = None) -> torch.nn.Module:
    """Move a torch model to target device and switch to eval mode.

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

    Returns:
        Inference-ready model moved to the resolved device.
    """
    return common_prepare_torch_model_for_inference(model, device)
