"""TensorRT engine validator for exported ``.engine`` 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 TensorRTValidator(BaseValidator):
    """Validate an exported TensorRT engine.

    Loads the serialised ``.engine`` file, allocates device/host buffers,
    runs inference on the provided *sample_input*, compares the output
    against an optional *baseline*, and benchmarks median inference latency.

    .. note::
        This validator **requires a CUDA-capable GPU** with both
        ``tensorrt`` and ``pycuda`` installed.
    """

    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 a TensorRT ``.engine`` file.

        Parameters
        ----------
        model_path:
            Path to the serialised TensorRT engine (``.engine``).
        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 TensorRT FP16/INT8 engines have larger numerical drift.
        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

        # ------------------------------------------------------------------
        # Lazy-import TensorRT and PyCUDA
        # ------------------------------------------------------------------
        try:
            import tensorrt as trt
        except ImportError:
            logger.warning(
                "tensorrt is not installed -- skipping TensorRT validation."
            )
            return {
                "status": "skipped",
                "reason": "tensorrt is not installed",
            }

        try:
            import pycuda.autoinit  # noqa: F401  (initialises CUDA context)
            import pycuda.driver as cuda
        except ImportError:
            logger.warning(
                "pycuda is not installed -- skipping TensorRT validation. "
                "Install it with:  pip install pycuda"
            )
            return {
                "status": "skipped",
                "reason": "pycuda is not installed",
            }

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

        # ------------------------------------------------------------------
        # 1. Deserialise the engine
        # ------------------------------------------------------------------
        try:
            trt_logger = trt.Logger(trt.Logger.WARNING)
            runtime = trt.Runtime(trt_logger)
            with open(model_path, "rb") as f:
                engine = runtime.deserialize_cuda_engine(f.read())
            if engine is None:
                return {
                    "status": "error",
                    "error": f"Failed to deserialise engine at {model_path}",
                }
            context = engine.create_execution_context()
        except Exception as exc:
            logger.error("TensorRTValidator: failed to load engine: %s", exc)
            return {"status": "error", "error": str(exc)}

        # ------------------------------------------------------------------
        # 2. Allocate host and device buffers
        # ------------------------------------------------------------------
        try:
            inputs_buf: list[dict[str, Any]] = []
            outputs_buf: list[dict[str, Any]] = []
            bindings: list[int] = []
            stream = cuda.Stream()

            for i in range(engine.num_bindings):
                name = engine.get_binding_name(i)
                dtype = trt.nptype(engine.get_binding_dtype(i))
                shape = engine.get_binding_shape(i)

                # Replace dynamic dimensions (-1) with sample_input shape
                resolved_shape = tuple(
                    sample_input.shape[j] if s == -1 else s
                    for j, s in enumerate(shape)
                )
                context.set_binding_shape(i, resolved_shape)

                size = int(np.prod(resolved_shape))
                host_mem = cuda.pagelocked_empty(size, dtype)
                device_mem = cuda.mem_alloc(host_mem.nbytes)
                bindings.append(int(device_mem))

                entry = {
                    "name": name,
                    "host": host_mem,
                    "device": device_mem,
                    "shape": resolved_shape,
                    "dtype": dtype,
                }
                if engine.binding_is_input(i):
                    inputs_buf.append(entry)
                else:
                    outputs_buf.append(entry)

        except Exception as exc:
            logger.error("TensorRTValidator: buffer allocation failed: %s", exc)
            return {"status": "error", "error": str(exc)}

        # ------------------------------------------------------------------
        # 3. Helper to run a single inference
        # ------------------------------------------------------------------
        def _infer() -> np.ndarray:
            # Copy input to host buffer, then host -> device
            np.copyto(
                inputs_buf[0]["host"],
                sample_input.astype(inputs_buf[0]["dtype"]).ravel(),
            )
            cuda.memcpy_htod_async(
                inputs_buf[0]["device"],
                inputs_buf[0]["host"],
                stream,
            )
            # Execute
            context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
            # Copy outputs device -> host
            for out in outputs_buf:
                cuda.memcpy_dtoh_async(out["host"], out["device"], stream)
            stream.synchronize()
            return outputs_buf[0]["host"].reshape(outputs_buf[0]["shape"]).copy()

        # ------------------------------------------------------------------
        # 4. Run inference and compare
        # ------------------------------------------------------------------
        try:
            output = _infer()
        except Exception as exc:
            logger.error("TensorRTValidator: inference failed: %s", exc)
            return {"status": "error", "error": str(exc)}

        comparison = self._compare_outputs(output, baseline, atol=atol, rtol=rtol)

        # ------------------------------------------------------------------
        # 5. Benchmark latency
        # ------------------------------------------------------------------
        def _run() -> None:
            _infer()

        latency_ms = self._benchmark_latency(_run)

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