"""OpenVINO validator for exported OpenVINO IR models."""

from __future__ import annotations

import logging
from pathlib import Path
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 OpenVinoValidator(BaseValidator):
    """Validate an exported OpenVINO IR model.

    Loads the model with ``openvino.runtime.Core``, compiles it for the
    CPU device, 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-3,
        rtol: float = 1e-3,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Validate an OpenVINO IR model.

        Parameters
        ----------
        model_path:
            Path to the OpenVINO model directory (containing ``.xml`` and
            ``.bin`` files) **or** directly to the ``.xml`` file.
        sample_input:
            Numpy array (typically ``float32``, NCHW layout).
        baseline:
            Optional reference output from the original model.
        atol:
            Absolute tolerance for value comparison.  Defaults to ``1e-3``
            because OpenVINO FP16 compression can introduce larger drift.
        rtol:
            Relative tolerance for value comparison.
        **kwargs:
            Extra arguments.  Recognised keys:

            * ``device`` (str) -- OpenVINO device name (default ``"CPU"``).

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

        try:
            from openvino.runtime import Core
        except ImportError:
            logger.warning(
                "openvino is not installed -- skipping OpenVINO validation. "
                "Install it with:  pip install openvino>=2023.0"
            )
            return {
                "status": "skipped",
                "reason": "openvino is not installed",
            }

        device: str = kwargs.get("device", "CPU")

        logger.info("OpenVinoValidator: loading %s (device=%s)", model_path, device)

        # ------------------------------------------------------------------
        # 1. Resolve the .xml path
        # ------------------------------------------------------------------
        model_file = Path(model_path)
        if model_file.is_dir():
            # Convention: look for a .xml file inside the directory
            xml_files = list(model_file.glob("*.xml"))
            if not xml_files:
                msg = f"No .xml file found in OpenVINO model directory: {model_file}"
                logger.error("OpenVinoValidator: %s", msg)
                return {"status": "error", "error": msg}
            xml_path = str(xml_files[0])
        else:
            xml_path = str(model_file)

        # ------------------------------------------------------------------
        # 2. Load and compile
        # ------------------------------------------------------------------
        try:
            core = Core()
            ov_model = core.read_model(xml_path)
            compiled = core.compile_model(ov_model, device)
        except Exception as exc:
            logger.error("OpenVinoValidator: failed to load/compile model: %s", exc)
            return {"status": "error", "error": str(exc)}

        # ------------------------------------------------------------------
        # 3. Run inference
        # ------------------------------------------------------------------
        try:
            input_layer = compiled.input(0)
            output_layer = compiled.output(0)

            # Ensure the input is float32
            input_data = sample_input.astype(np.float32)

            infer_request = compiled.create_infer_request()
            result = infer_request.infer({input_layer: input_data})
            output = result[output_layer]
        except Exception as exc:
            logger.error("OpenVinoValidator: inference failed: %s", exc)
            return {"status": "error", "error": str(exc)}

        # ------------------------------------------------------------------
        # 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:
            infer_request.infer({input_layer: input_data})

        latency_ms = self._benchmark_latency(_run)

        logger.info(
            "OpenVinoValidator: 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,
        }
