"""ONNX model exporter."""

from __future__ import annotations

import logging
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 OnnxExporter(BaseExporter):
    """Export a PyTorch model to the ONNX format.

    Uses ``torch.onnx.export`` with configurable opset version, dynamic axes,
    metadata embedding (stride, class names), and optional graph simplification
    via ``onnx-simplifier``.
    """

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

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

    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 ONNX.

        Args:
            model: PyTorch model in eval mode.
            sample_input: Sample input tensor (BCHW).
            output_dir: Directory to write the exported file into.
            file_stem: Base filename without extension.
            **kwargs:
                opset (int): ONNX opset version.  Defaults to ``12``.
                dynamic (bool): Enable dynamic axes for batch / spatial dims.
                    Defaults to ``False``.
                simplify (bool): Run ``onnx-simplifier`` after export.
                    Defaults to ``False``.
                stride (int | None): Model stride value to embed in metadata.
                names (dict | list | None): Class names to embed in metadata.
                output_names (list[str] | None): Custom output tensor names.
                    Detected automatically when the model exposes a
                    ``SegmentationModel`` or ``DetectionModel`` type.
                is_segmentation (bool): Hint that the model is a segmentation
                    model so the correct dynamic axes are generated.
                is_detection (bool): Hint that the model is a detection model.

        Returns:
            Absolute path to the exported ``.onnx`` file.
        """
        import torch

        try:
            import onnx
        except ImportError as exc:
            raise ImportError(
                "The 'onnx' package (>= 1.12.0) is required for ONNX export. "
                "Install it with:  pip install onnx>=1.12.0"
            ) from exc

        opset: int = kwargs.get("opset", 12)
        dynamic: bool = kwargs.get("dynamic", False)
        simplify: bool = kwargs.get("simplify", False)
        stride: int | None = kwargs.get("stride", None)
        names = kwargs.get("names", None)
        output_names: list[str] | None = kwargs.get("output_names", None)
        is_segmentation: bool = kwargs.get("is_segmentation", False)
        is_detection: bool = kwargs.get("is_detection", False)

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

        # ------------------------------------------------------------------
        # Determine output names automatically when not provided
        # ------------------------------------------------------------------
        if output_names is None:
            output_names = ["output0", "output1"] if is_segmentation else ["output0"]

        # ------------------------------------------------------------------
        # Dynamic axes
        # ------------------------------------------------------------------
        dynamic_axes: dict | None = None
        if dynamic:
            dynamic_axes = {"images": {0: "batch", 2: "height", 3: "width"}}
            if is_segmentation:
                dynamic_axes["output0"] = {0: "batch", 1: "anchors"}
                dynamic_axes["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}
            elif is_detection:
                dynamic_axes["output0"] = {0: "batch", 1: "anchors"}

        logger.info(
            "ONNX export: starting with onnx %s, opset %d ...",
            onnx.__version__,
            opset,
        )

        # ------------------------------------------------------------------
        # torch.onnx.export
        # ------------------------------------------------------------------
        export_model = model.cpu() if dynamic else model
        export_input = sample_input.cpu() if dynamic else sample_input

        torch.onnx.export(
            export_model,
            export_input,
            f,
            verbose=False,
            opset_version=opset,
            do_constant_folding=True,
            input_names=["images"],
            output_names=output_names,
            dynamic_axes=dynamic_axes,
        )

        # ------------------------------------------------------------------
        # Validate
        # ------------------------------------------------------------------
        model_onnx = onnx.load(f)
        onnx.checker.check_model(model_onnx)
        logger.info("ONNX export: model validation passed.")

        # ------------------------------------------------------------------
        # Embed metadata
        # ------------------------------------------------------------------
        computed_stride = (
            stride
            if stride is not None
            else int(max(model.stride)) if hasattr(model, "stride") else 1
        )
        metadata: dict = {"stride": computed_stride, "names": names}
        for k, v in metadata.items():
            meta = model_onnx.metadata_props.add()
            meta.key, meta.value = k, str(v)
        onnx.save(model_onnx, f)

        # ------------------------------------------------------------------
        # Optional: simplify
        # ------------------------------------------------------------------
        if simplify:
            try:
                import onnxsim
            except ImportError as exc:
                raise ImportError(
                    "onnx-simplifier (>= 0.4.1) is required when simplify=True. "
                    "Install it with:  pip install onnx-simplifier>=0.4.1"
                ) from exc

            logger.info(
                "ONNX export: simplifying with onnx-simplifier %s ...",
                onnxsim.__version__,
            )
            model_onnx, check = onnxsim.simplify(model_onnx)
            if not check:
                raise RuntimeError(
                    "onnx-simplifier validation failed after simplification."
                )
            onnx.save(model_onnx, f)

        logger.info("ONNX export: saved %s", output_path)
        return f
