"""Adapter for converting an ONNX model to downstream runtime formats."""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class OnnxInputAdapter:
    """Convert an existing ONNX model to downstream inference formats.

    Each conversion method lazily imports the required third-party library
    so that users only need the dependencies for the formats they actually
    target.  Missing libraries produce a clear :class:`ImportError` with
    install instructions rather than a cryptic traceback.

    Supported downstream formats (from an ``.onnx`` starting point):

    * **openvino** -- via ``openvino`` (``openvino-dev`` / ``openvino>=2023.0``)
    * **engine** (TensorRT) -- via ``tensorrt``
    * **coreml** -- via ``coremltools`` + ``onnx``
    * **tflite** -- via ``onnx-tf`` + ``tensorflow``

    Example
    -------
    ::

        adapter = OnnxInputAdapter()
        ov_path = adapter.export("model.onnx", "openvino", "./exports")
    """

    SUPPORTED_FORMATS: list[str] = ["openvino", "engine", "coreml", "tflite"]

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #
    def export(
        self,
        onnx_path: str,
        fmt: str,
        output_dir: str,
        **kwargs: Any,
    ) -> str:
        """Convert an ONNX model to the target *fmt*.

        Parameters
        ----------
        onnx_path:
            Path to the source ``.onnx`` file.
        fmt:
            Target format name.  Must be one of :attr:`SUPPORTED_FORMATS`.
        output_dir:
            Directory where the converted artefact will be written.
        **kwargs:
            Format-specific keyword arguments (e.g. ``half``, ``workspace``).

        Returns
        -------
        str
            Path to the converted artefact (file or directory).

        Raises
        ------
        FileNotFoundError
            If *onnx_path* does not exist.
        ValueError
            If *fmt* is not a supported format.
        ImportError
            If the required third-party library is not installed.
        """
        onnx_path = str(onnx_path)
        if not os.path.isfile(onnx_path):
            raise FileNotFoundError(f"ONNX model not found: {onnx_path}")

        os.makedirs(output_dir, exist_ok=True)

        dispatch = {
            "openvino": self.to_openvino,
            "engine": self.to_tensorrt,
            "coreml": self.to_coreml,
            "tflite": self.to_tflite,
        }

        handler = dispatch.get(fmt)
        if handler is None:
            raise ValueError(
                f"Unsupported ONNX conversion format: {fmt!r}. "
                f"Supported: {self.SUPPORTED_FORMATS}"
            )

        return handler(onnx_path, output_dir, **kwargs)

    # ------------------------------------------------------------------ #
    # OpenVINO
    # ------------------------------------------------------------------ #
    def to_openvino(
        self,
        onnx_path: str,
        output_dir: str,
        half: bool = False,
        **kwargs: Any,
    ) -> str:
        """Convert ONNX to OpenVINO IR (Intermediate Representation).

        Requires ``openvino>=2023.0`` (``pip install openvino``).

        Parameters
        ----------
        onnx_path:
            Path to the source ``.onnx`` file.
        output_dir:
            Directory for the output ``*_openvino_model/`` folder.
        half:
            If ``True``, compress weights to FP16.

        Returns
        -------
        str
            Path to the output OpenVINO model directory.
        """
        try:
            import openvino as ov
        except ImportError as exc:
            raise ImportError(
                "OpenVINO conversion requires the 'openvino' package. "
                "Install it with:  pip install openvino>=2023.0"
            ) from exc

        stem = Path(onnx_path).stem
        ov_dir = os.path.join(output_dir, f"{stem}_openvino_model")
        os.makedirs(ov_dir, exist_ok=True)

        logger.info(
            "Converting ONNX -> OpenVINO (half=%s): %s -> %s",
            half, onnx_path, ov_dir,
        )

        ov_model = ov.convert_model(onnx_path)

        xml_path = os.path.join(ov_dir, f"{stem}.xml")
        ov.save_model(ov_model, xml_path, compress_to_fp16=half)

        logger.info("OpenVINO model saved to %s", ov_dir)
        return ov_dir

    # ------------------------------------------------------------------ #
    # TensorRT
    # ------------------------------------------------------------------ #
    def to_tensorrt(
        self,
        onnx_path: str,
        output_dir: str,
        half: bool = False,
        dynamic: bool = False,
        workspace: int = 4,
        max_batch: int = 16,
        opt_batch: int | None = None,
        **kwargs: Any,
    ) -> str:
        """Convert ONNX to a TensorRT serialised engine.

        Requires ``tensorrt`` (``pip install nvidia-tensorrt`` or the
        NVIDIA TensorRT distribution).

        Parameters
        ----------
        onnx_path:
            Path to the source ``.onnx`` file.
        output_dir:
            Directory for the output ``.engine`` file.
        half:
            If ``True``, build the engine with FP16 precision.
        dynamic:
            If ``True``, create an optimization profile that allows
            variable batch sizes at inference time (min=1, opt, max).
        workspace:
            Workspace size in GiB for the TensorRT builder.
        max_batch:
            Maximum batch size for the dynamic optimization profile.
            Only used when ``dynamic=True``.  Defaults to ``16``.
        opt_batch:
            Optimal batch size for the dynamic optimization profile.
            Only used when ``dynamic=True``.  Defaults to ``max_batch // 2``.

        Returns
        -------
        str
            Path to the serialised ``.engine`` file.
        """
        try:
            import tensorrt as trt
        except ImportError as exc:
            raise ImportError(
                "TensorRT conversion requires the 'tensorrt' package. "
                "Install it with:  pip install nvidia-tensorrt  "
                "(or use the NVIDIA TensorRT distribution for your platform)."
            ) from exc

        stem = Path(onnx_path).stem
        engine_path = os.path.join(output_dir, f"{stem}.engine")

        logger.info(
            "Converting ONNX -> TensorRT (half=%s, dynamic=%s, workspace=%dGiB): %s -> %s",
            half, dynamic, workspace, onnx_path, engine_path,
        )

        trt_logger = trt.Logger(trt.Logger.INFO)
        builder = trt.Builder(trt_logger)
        config = builder.create_builder_config()

        # TensorRT 8.4+ uses set_memory_pool_limit instead of max_workspace_size
        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)

        # When dynamic=True, make the ONNX batch dimension symbolic so
        # TensorRT allows an optimization profile with variable batch.
        onnx_bytes = Path(onnx_path).read_bytes()
        if dynamic:
            try:
                import onnx as _onnx

                onnx_model = _onnx.load_from_string(onnx_bytes)
                for inp in onnx_model.graph.input:
                    if inp.type.tensor_type.shape.dim:
                        inp.type.tensor_type.shape.dim[0].dim_param = "batch"
                onnx_bytes = onnx_model.SerializeToString()
                logger.info("Patched ONNX input batch dimension to dynamic.")
            except ImportError:
                logger.warning(
                    "onnx package not installed; cannot make batch dim dynamic. "
                    "The engine may fail to build with dynamic=True."
                )

        if not parser.parse(onnx_bytes):
            errors = "\n".join(
                str(parser.get_error(i)) for i in range(parser.num_errors)
            )
            raise RuntimeError(
                f"Failed to parse ONNX model:\n{errors}"
            )

        # Dynamic batch optimization profile
        if dynamic:
            effective_opt = opt_batch if opt_batch is not None else max(1, max_batch // 2)
            profile = builder.create_optimization_profile()
            for i in range(network.num_inputs):
                inp = network.get_input(i)
                shape = inp.shape  # e.g. (1, C, H, W) or (-1, C, H, W)
                spatial = tuple(
                    s if s > 0 else 1 for s in shape[1:]
                )
                profile.set_shape(
                    inp.name,
                    (1, *spatial),              # min
                    (effective_opt, *spatial),   # opt
                    (max_batch, *spatial),       # max
                )
                logger.info(
                    "Dynamic profile for '%s': min=(1,%s) opt=(%d,%s) max=(%d,%s)",
                    inp.name,
                    ",".join(str(s) for s in spatial),
                    effective_opt,
                    ",".join(str(s) for s in spatial),
                    max_batch,
                    ",".join(str(s) for s in spatial),
                )
            config.add_optimization_profile(profile)

        if half and builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            logger.info("FP16 mode enabled for TensorRT engine.")
        elif half:
            logger.warning("FP16 requested but platform does not support fast FP16.")

        serialized = builder.build_serialized_network(network, config)
        if serialized is None:
            raise RuntimeError("TensorRT engine build returned None -- check logs.")

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

        logger.info("TensorRT engine saved to %s", engine_path)
        return engine_path

    # ------------------------------------------------------------------ #
    # CoreML
    # ------------------------------------------------------------------ #
    def to_coreml(
        self,
        onnx_path: str,
        output_dir: str,
        **kwargs: Any,
    ) -> str:
        """Convert ONNX to a CoreML ``.mlpackage``.

        Requires ``coremltools`` and ``onnx``
        (``pip install coremltools onnx``).

        Parameters
        ----------
        onnx_path:
            Path to the source ``.onnx`` file.
        output_dir:
            Directory for the output ``.mlpackage``.

        Returns
        -------
        str
            Path to the saved ``.mlpackage`` directory.
        """
        try:
            import coremltools as ct
        except ImportError as exc:
            raise ImportError(
                "CoreML conversion requires the 'coremltools' package. "
                "Install it with:  pip install coremltools"
            ) from exc

        try:
            import onnx  # noqa: F811 -- intentional lazy import
        except ImportError as exc:
            raise ImportError(
                "CoreML conversion from ONNX also requires the 'onnx' package. "
                "Install it with:  pip install onnx"
            ) from exc

        stem = Path(onnx_path).stem
        mlpackage_path = os.path.join(output_dir, f"{stem}.mlpackage")

        logger.info(
            "Converting ONNX -> CoreML: %s -> %s",
            onnx_path, mlpackage_path,
        )

        onnx_model = onnx.load(onnx_path)
        ct_model = ct.converters.onnx.convert(model=onnx_model)
        ct_model.save(mlpackage_path)

        logger.info("CoreML model saved to %s", mlpackage_path)
        return mlpackage_path

    # ------------------------------------------------------------------ #
    # TFLite
    # ------------------------------------------------------------------ #
    def to_tflite(
        self,
        onnx_path: str,
        output_dir: str,
        **kwargs: Any,
    ) -> str:
        """Convert ONNX to TensorFlow Lite (``.tflite``).

        This uses a two-step process:

        1. ONNX -> TensorFlow SavedModel  (via ``onnx-tf``)
        2. SavedModel -> TFLite           (via ``tensorflow``)

        Requires ``onnx-tf`` and ``tensorflow``
        (``pip install onnx-tf tensorflow``).

        Parameters
        ----------
        onnx_path:
            Path to the source ``.onnx`` file.
        output_dir:
            Directory for the output ``.tflite`` file.

        Returns
        -------
        str
            Path to the converted ``.tflite`` file.
        """
        try:
            import onnx  # noqa: F811
        except ImportError as exc:
            raise ImportError(
                "TFLite conversion requires the 'onnx' package. "
                "Install it with:  pip install onnx"
            ) from exc

        try:
            from onnx_tf.backend import prepare  # type: ignore[import-untyped]
        except ImportError as exc:
            raise ImportError(
                "TFLite conversion requires the 'onnx-tf' package. "
                "Install it with:  pip install onnx-tf"
            ) from exc

        try:
            import tensorflow as tf
        except ImportError as exc:
            raise ImportError(
                "TFLite conversion requires 'tensorflow'. "
                "Install it with:  pip install tensorflow"
            ) from exc

        stem = Path(onnx_path).stem
        saved_model_dir = os.path.join(output_dir, f"{stem}_saved_model")
        tflite_path = os.path.join(output_dir, f"{stem}.tflite")

        logger.info(
            "Converting ONNX -> TFLite (via SavedModel): %s -> %s",
            onnx_path, tflite_path,
        )

        # Step 1: ONNX -> TF SavedModel
        onnx_model = onnx.load(onnx_path)
        tf_rep = prepare(onnx_model)
        tf_rep.export_graph(saved_model_dir)

        # Step 2: SavedModel -> TFLite
        converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
        converter.optimizations = [tf.lite.Optimize.DEFAULT]

        tflite_model = converter.convert()
        with open(tflite_path, "wb") as f:
            f.write(tflite_model)

        logger.info("TFLite model saved to %s", tflite_path)
        return tflite_path
