"""ONNX Runtime validator for exported ``.onnx`` models."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import numpy as np

from matrice_export.validators.base import BaseValidator

logger = logging.getLogger(__name__)


class OnnxValidator(BaseValidator):
    """Validate an exported ONNX model using ONNX Runtime.

    Loads the model with ``onnxruntime.InferenceSession``, runs inference on
    the provided *sample_input*, compares the output against an optional
    *baseline*, and benchmarks median inference latency.
    """

    def validate(
        self,
        model_path: str,
        sample_input: np.ndarray,
        baseline: np.ndarray | None = None,
        atol: float = 1e-5,
        rtol: float = 1e-3,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Validate an ``.onnx`` model.

        Parameters
        ----------
        model_path:
            Path to the ``.onnx`` file.
        sample_input:
            Numpy array (typically ``float32``, NCHW layout).
        baseline:
            Optional reference output from the original model.
        atol:
            Absolute tolerance for value comparison.
        rtol:
            Relative tolerance for value comparison.
        **kwargs:
            Extra arguments (ignored).

        Returns
        -------
        dict
            Validation result with keys ``shape_match``, ``values_match``,
            ``max_diff``, ``latency_ms``, ``output_shape``.
        """
        import numpy as np

        try:
            import onnxruntime as ort
        except ImportError:
            logger.warning(
                "onnxruntime is not installed -- skipping ONNX validation. "
                "Install it with:  pip install onnxruntime"
            )
            return {
                "status": "skipped",
                "reason": "onnxruntime is not installed",
            }

        logger.info("OnnxValidator: loading %s", model_path)

        # ------------------------------------------------------------------
        # 1. Create inference session
        # ------------------------------------------------------------------
        try:
            providers = ort.get_available_providers()
            # Prefer CUDA if available, but always fall back to CPU
            preferred = []
            if "CUDAExecutionProvider" in providers:
                preferred.append("CUDAExecutionProvider")
            preferred.append("CPUExecutionProvider")

            sess = ort.InferenceSession(model_path, providers=preferred)
        except Exception as exc:
            logger.error("OnnxValidator: failed to load model: %s", exc)
            return {"status": "error", "error": str(exc)}

        input_name = sess.get_inputs()[0].name

        # Ensure sample_input dtype matches the model's expected dtype
        input_feed = {input_name: sample_input.astype(np.float32)}

        # ------------------------------------------------------------------
        # 2. Run inference
        # ------------------------------------------------------------------
        try:
            outputs = sess.run(None, input_feed)
        except Exception as exc:
            logger.error("OnnxValidator: inference failed: %s", exc)
            return {"status": "error", "error": str(exc)}

        output = outputs[0]

        # ------------------------------------------------------------------
        # 3 & 4. Compare shapes and values against baseline
        # ------------------------------------------------------------------
        comparison = self._compare_outputs(output, baseline, atol=atol, rtol=rtol)

        # ------------------------------------------------------------------
        # 5. Benchmark latency (median of 100 runs)
        # ------------------------------------------------------------------
        def _run() -> None:
            sess.run(None, input_feed)

        latency_ms = self._benchmark_latency(_run)

        logger.info(
            "OnnxValidator: shape_match=%s  values_match=%s  max_diff=%s  latency=%.2f ms",
            comparison["shape_match"],
            comparison["values_match"],
            comparison["max_diff"],
            latency_ms,
        )

        return {
            **comparison,
            "latency_ms": latency_ms,
        }
