"""Windows iGPU inference via DirectML + ONNX Runtime.

DirectML is Microsoft's hardware-accelerated machine-learning backend that
runs on any DirectX 12 capable GPU, including integrated GPUs on Windows
laptops.  This module wraps ``onnxruntime`` with the ``DmlExecutionProvider``
and provides helpers for loading, inference, device introspection, and
benchmarking.
"""

from __future__ import annotations

import logging
import time
from pathlib import Path
from typing import Any, Union

logger = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Load
# --------------------------------------------------------------------------- #

def load_onnx_directml(
    model_path: Union[str, Path],
    device_id: int = 0,
) -> Any:
    """Load an ONNX model with the DirectML execution provider.

    Falls back to ``CPUExecutionProvider`` if DirectML is not available
    (e.g. on Linux / macOS or when the ``onnxruntime-directml`` package is
    not installed).

    Args:
        model_path: Path to an ``.onnx`` file.
        device_id: DirectML device index (0 = default GPU).

    Returns:
        An ``onnxruntime.InferenceSession`` configured with the best
        available provider.

    Raises:
        ImportError: If ``onnxruntime`` is not installed.
        FileNotFoundError: If *model_path* does not exist.
    """
    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "DirectML inference requires 'onnxruntime-directml'. "
            "Install it with:  pip install onnxruntime-directml"
        ) from exc

    model_path = Path(model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"ONNX model not found: {model_path}")

    # Build provider list with DirectML preferred
    providers: list[Union[str, tuple[str, dict[str, Any]]]] = []
    available = ort.get_available_providers()

    if "DmlExecutionProvider" in available:
        providers.append(
            ("DmlExecutionProvider", {"device_id": device_id})
        )
        logger.info(
            "DirectML execution provider available; using device_id=%d.",
            device_id,
        )
    else:
        logger.info(
            "DmlExecutionProvider not available (providers: %s); "
            "falling back to CPU.",
            available,
        )

    providers.append("CPUExecutionProvider")

    # Session options
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = (
        ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    )
    sess_options.enable_mem_pattern = True

    logger.info("Loading ONNX model from %s ...", model_path)
    session = ort.InferenceSession(
        str(model_path),
        sess_options,
        providers=providers,
    )

    active_provider = session.get_providers()[0] if session.get_providers() else "unknown"
    logger.info("InferenceSession created with provider: %s", active_provider)

    return session


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #

def run_inference_directml(
    session: Any,
    input_data: Any,
    output_names: list[str] | None = None,
) -> list[Any]:
    """Run inference using a DirectML (or CPU fallback) session.

    Args:
        session: An ``onnxruntime.InferenceSession`` returned by
            :func:`load_onnx_directml`.
        input_data: Input as a ``numpy.ndarray``.  Must match the shape and
            dtype expected by the model's first input.
        output_names: Optional list of output tensor names to retrieve.
            Defaults to ``None`` (all outputs).

    Returns:
        A list of ``numpy.ndarray`` outputs, one per model output.
    """
    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError(
            "DirectML inference requires 'numpy'. "
            "Install it with:  pip install numpy"
        ) from exc

    input_meta = session.get_inputs()[0]
    input_name = input_meta.name

    # Ensure correct dtype
    if hasattr(input_data, "dtype"):
        expected_type = input_meta.type
        if "float16" in expected_type:
            input_data = np.asarray(input_data, dtype=np.float16)
        elif "float" in expected_type:
            input_data = np.asarray(input_data, dtype=np.float32)
        else:
            input_data = np.asarray(input_data)

    return session.run(output_names, {input_name: input_data})


# --------------------------------------------------------------------------- #
# Device info
# --------------------------------------------------------------------------- #

def get_directml_info() -> dict[str, Any]:
    """Return DirectML availability and device information.

    Returns:
        A dict with keys:

        * ``available`` -- ``True`` if ``DmlExecutionProvider`` is usable
        * ``provider`` -- the name of the preferred provider
        * ``all_providers`` -- list of all available ORT providers
        * ``onnxruntime_version`` -- installed onnxruntime version string
        * ``device_name`` -- human-readable device description (best-effort)
    """
    result: dict[str, Any] = {
        "available": False,
        "provider": "CPUExecutionProvider",
        "all_providers": [],
        "onnxruntime_version": None,
        "device_name": "CPU (fallback)",
    }

    try:
        import onnxruntime as ort
    except ImportError:
        logger.warning(
            "onnxruntime is not installed; DirectML info unavailable."
        )
        return result

    result["onnxruntime_version"] = ort.__version__
    available_providers = ort.get_available_providers()
    result["all_providers"] = available_providers

    if "DmlExecutionProvider" in available_providers:
        result["available"] = True
        result["provider"] = "DmlExecutionProvider"
        result["device_name"] = "DirectML GPU (DirectX 12)"
        logger.info("DirectML is available.")
    else:
        logger.info(
            "DirectML is NOT available.  Available providers: %s",
            available_providers,
        )

    return result


# --------------------------------------------------------------------------- #
# Benchmark
# --------------------------------------------------------------------------- #

def benchmark_directml(
    session: Any,
    input_data: Any,
    n_runs: int = 100,
    warmup_runs: int = 10,
) -> dict[str, Any]:
    """Benchmark inference latency with a DirectML (or CPU) session.

    Args:
        session: An ``onnxruntime.InferenceSession``.
        input_data: Input ``numpy.ndarray`` matching the model's expected
            shape and dtype.
        n_runs: Number of timed inference runs.
        warmup_runs: Number of warm-up runs (discarded).

    Returns:
        A dict with keys:

        * ``provider`` -- the active execution provider name
        * ``n_runs`` -- number of timed runs
        * ``latency_mean_ms`` -- mean latency in milliseconds
        * ``latency_median_ms`` -- median latency
        * ``latency_min_ms`` -- minimum latency
        * ``latency_max_ms`` -- maximum latency
        * ``throughput_fps`` -- estimated frames per second
    """
    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError(
            "DirectML benchmarking requires 'numpy'. "
            "Install it with:  pip install numpy"
        ) from exc

    input_name = session.get_inputs()[0].name
    provider = (
        session.get_providers()[0] if session.get_providers() else "unknown"
    )

    logger.info(
        "Benchmarking on %s: %d warmup + %d timed runs ...",
        provider,
        warmup_runs,
        n_runs,
    )

    # Warm-up
    for _ in range(warmup_runs):
        session.run(None, {input_name: input_data})

    # Timed runs
    latencies: list[float] = []
    for _ in range(n_runs):
        t_start = time.perf_counter()
        session.run(None, {input_name: input_data})
        t_end = time.perf_counter()
        latencies.append((t_end - t_start) * 1000.0)

    latencies_np = np.array(latencies)
    mean_ms = float(np.mean(latencies_np))

    result: dict[str, Any] = {
        "provider": provider,
        "n_runs": n_runs,
        "latency_mean_ms": mean_ms,
        "latency_median_ms": float(np.median(latencies_np)),
        "latency_min_ms": float(np.min(latencies_np)),
        "latency_max_ms": float(np.max(latencies_np)),
        "throughput_fps": 1000.0 / mean_ms if mean_ms > 0 else 0.0,
    }

    logger.info(
        "Benchmark results (%s): mean=%.2f ms, median=%.2f ms, "
        "min=%.2f ms, max=%.2f ms, throughput=%.1f FPS",
        provider,
        result["latency_mean_ms"],
        result["latency_median_ms"],
        result["latency_min_ms"],
        result["latency_max_ms"],
        result["throughput_fps"],
    )

    return result
