"""TensorRT engine exporter."""

from __future__ import annotations

import logging
import platform
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import torch

from matrice_export.formats.base import BaseExporter

logger = logging.getLogger(__name__)


class TensorRTExporter(BaseExporter):
    """Export a PyTorch model to a TensorRT serialised engine.

    The pipeline works in two stages:

    1. The model is first exported to ONNX (using
       :class:`~matrice_export.formats.onnx_export.OnnxExporter` or a
       pre-existing ONNX file).
    2. A TensorRT engine is built from the ONNX graph using
       ``tensorrt.Builder``, supporting dynamic shapes and FP16
       precision.

    .. note::
        This exporter **requires a GPU** (``requires_gpu = True``).
    """

    @property
    def format_name(self) -> str:
        return "tensorrt"

    @property
    def suffix(self) -> str:
        return ".engine"

    @property
    def requires_gpu(self) -> bool:
        return True

    def export(
        self,
        model: torch.nn.Module,
        sample_input: torch.Tensor,
        output_dir: str | Path,
        file_stem: str = "model",
        **kwargs,
    ) -> str:
        """Export a PyTorch model to a TensorRT engine.

        Args:
            model: PyTorch model in eval mode.
            sample_input: Sample input tensor (BCHW) **on a CUDA device**.
            output_dir: Directory to write the exported engine into.
            file_stem: Base filename without extension.
            **kwargs:
                half (bool): Build an FP16 engine when the platform supports
                    fast FP16.  Defaults to ``False``.
                dynamic (bool): Enable dynamic batch / spatial dimensions.
                    Defaults to ``False``.
                simplify (bool): Simplify the intermediate ONNX graph.
                    Defaults to ``False``.
                workspace (int): Maximum workspace size in GiB for the TRT
                    builder.  Defaults to ``4``.
                verbose (bool): Enable verbose TRT builder logging.
                    Defaults to ``False``.
                opset (int): ONNX opset version for the intermediate export.
                    Defaults to ``12``.
                onnx_path (str | Path | None): Path to a pre-existing ONNX
                    model.  When ``None`` the exporter generates one.
                max_batch (int): Maximum batch size for the dynamic
                    optimization profile.  Defaults to ``0`` (auto —
                    uses ``max(sample_input.shape[0], 16)``).
                opt_batch (int | None): Optimal batch size for the
                    dynamic optimization profile.  Defaults to
                    ``max_batch // 2``.
                stride (int | None): Model stride value for ONNX metadata.
                names (dict | list | None): Class names for ONNX metadata.

        Returns:
            Absolute path to the exported ``.engine`` file.

        Raises:
            RuntimeError: If no CUDA device is available or the ONNX file
                cannot be parsed.
        """
        import torch

        try:
            import tensorrt as trt
        except ImportError:
            if platform.system() == "Linux":
                raise ImportError(
                    "The 'tensorrt' package is required for TensorRT export. "
                    "Install it with:  pip install nvidia-tensorrt "
                    "-U --index-url https://pypi.ngc.nvidia.com"
                )
            raise ImportError(
                "The 'tensorrt' package is required for TensorRT export. "
                "Please install the appropriate TensorRT package for your platform."
            )

        half: bool = kwargs.get("half", False)
        dynamic: bool = kwargs.get("dynamic", False)
        simplify: bool = kwargs.get("simplify", False)
        workspace: int = kwargs.get("workspace", 4)
        verbose: bool = kwargs.get("verbose", False)
        opset: int = kwargs.get("opset", 12)
        onnx_path: str | Path | None = kwargs.get("onnx_path", None)
        max_batch: int = kwargs.get("max_batch", 0)  # 0 = auto
        opt_batch: int | None = kwargs.get("opt_batch", None)

        if sample_input.device.type == "cpu":
            raise RuntimeError(
                "TensorRT export requires a CUDA device but sample_input "
                "is on CPU.  Move both model and input to a GPU first."
            )

        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        engine_path = output_dir / f"{file_stem}{self.suffix}"

        # ------------------------------------------------------------------
        # Stage 1 -- produce an ONNX intermediate
        # ------------------------------------------------------------------
        if onnx_path is None:
            onnx_path = output_dir / f"{file_stem}.onnx"
            if not Path(onnx_path).exists():
                logger.info(
                    "TensorRT export: ONNX prerequisite not found, "
                    "running OnnxExporter first ..."
                )
                from matrice_export.formats.onnx_export import OnnxExporter

                onnx_exporter = OnnxExporter()
                onnx_path = onnx_exporter.export(
                    model,
                    sample_input,
                    output_dir,
                    file_stem=file_stem,
                    opset=opset,
                    dynamic=dynamic,
                    simplify=simplify,
                    **{k: v for k, v in kwargs.items() if k in ("stride", "names", "is_segmentation", "is_detection")},
                )

        onnx_path = Path(onnx_path)
        if not onnx_path.exists():
            raise FileNotFoundError(
                f"TensorRT export: expected ONNX file at {onnx_path} but it "
                "does not exist."
            )

        logger.info("TensorRT export: building engine from %s ...", onnx_path)

        # ------------------------------------------------------------------
        # Stage 2 -- build the TensorRT engine
        # ------------------------------------------------------------------
        trt_logger = trt.Logger(trt.Logger.INFO)
        if verbose:
            trt_logger.min_severity = trt.Logger.Severity.VERBOSE

        builder = trt.Builder(trt_logger)
        config = builder.create_builder_config()
        # TensorRT 8.4+ deprecates max_workspace_size in favour of
        # set_memory_pool_limit.  Use the new API when available.
        if hasattr(config, "set_memory_pool_limit"):
            config.set_memory_pool_limit(
                trt.MemoryPoolType.WORKSPACE, workspace << 30
            )
        else:
            config.max_workspace_size = workspace << 30  # type: ignore[attr-defined]

        flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        network = builder.create_network(flag)
        parser = trt.OnnxParser(network, trt_logger)

        if not parser.parse_from_file(str(onnx_path)):
            errors = "\n".join(
                str(parser.get_error(i)) for i in range(parser.num_errors)
            )
            raise RuntimeError(
                f"TensorRT export: failed to parse ONNX file {onnx_path}.\n{errors}"
            )

        inputs = [network.get_input(i) for i in range(network.num_inputs)]
        outputs = [network.get_output(i) for i in range(network.num_outputs)]

        for inp in inputs:
            logger.info(
                "TensorRT export: input '%s' shape=%s dtype=%s",
                inp.name, inp.shape, inp.dtype,
            )
        for out in outputs:
            logger.info(
                "TensorRT export: output '%s' shape=%s dtype=%s",
                out.name, out.shape, out.dtype,
            )

        # ------------------------------------------------------------------
        # Dynamic shapes
        # ------------------------------------------------------------------
        if dynamic:
            # Determine effective max and opt batch sizes.
            # When max_batch=0 (auto), use at least 16 or the sample batch.
            effective_max = max_batch if max_batch > 0 else max(sample_input.shape[0], 16)
            effective_opt = opt_batch if opt_batch is not None else max(1, effective_max // 2)

            logger.info(
                "TensorRT export: dynamic batch profile min=1, opt=%d, max=%d",
                effective_opt, effective_max,
            )

            profile = builder.create_optimization_profile()
            for inp in inputs:
                profile.set_shape(
                    inp.name,
                    (1, *sample_input.shape[1:]),              # min
                    (effective_opt, *sample_input.shape[1:]),   # opt
                    (effective_max, *sample_input.shape[1:]),   # max
                )
            config.add_optimization_profile(profile)

        # ------------------------------------------------------------------
        # FP16
        # ------------------------------------------------------------------
        use_fp16 = builder.platform_has_fast_fp16 and half
        if use_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
        logger.info(
            "TensorRT export: building FP%d engine as %s",
            16 if use_fp16 else 32,
            engine_path,
        )

        # ------------------------------------------------------------------
        # Serialize
        # ------------------------------------------------------------------
        serialized = builder.build_serialized_network(network, config)
        if serialized is None:
            raise RuntimeError(
                "TensorRT export: builder.build_serialized_network returned None."
            )

        with open(engine_path, "wb") as f:
            f.write(bytearray(serialized))

        logger.info("TensorRT export: saved %s", engine_path)
        return str(engine_path)
