"""TensorFlow SavedModel exporter.

Converts a PyTorch model to TensorFlow SavedModel format using the portable
ONNX-to-TF pipeline (PyTorch -> ONNX -> TF SavedModel) via ``onnx-tf``.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import torch

from matrice_export.formats.base import BaseExporter

logger = logging.getLogger(__name__)


class TFSavedModelExporter(BaseExporter):
    """Export a PyTorch model to TensorFlow SavedModel via ONNX intermediate."""

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

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

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def export(
        self,
        model: torch.nn.Module,
        sample_input: torch.Tensor,
        output_dir: str | Path,
        file_stem: str = "model",
        **kwargs: Any,
    ) -> str:
        """Export PyTorch model to TF SavedModel.

        The portable approach first exports to ONNX, then converts the ONNX
        graph to a TensorFlow SavedModel using the ``onnx-tf`` package.

        Keyword Args:
            onnx_path (str | Path | None): Path to a pre-existing ONNX file.
                If *None*, the model is exported to ONNX automatically.
            opset (int): ONNX opset version (default ``12``).
            dynamic (bool): Whether to use dynamic axes in the ONNX export.

        Returns:
            Path to the exported SavedModel directory.
        """
        import torch  # noqa: F811 — lazy import

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

        onnx_path: str | Path | None = kwargs.get("onnx_path")
        opset: int = kwargs.get("opset", 12)
        dynamic: bool = kwargs.get("dynamic", False)

        # Step 1 -- obtain an ONNX file --------------------------------
        if onnx_path is None:
            onnx_path = self._export_onnx(
                model, sample_input, output_dir, file_stem, opset, dynamic
            )
        onnx_path = Path(onnx_path)
        if not onnx_path.exists():
            raise FileNotFoundError(f"ONNX file not found: {onnx_path}")

        # Step 2 -- convert ONNX -> TF SavedModel ----------------------
        try:
            import onnx
            from onnx_tf.backend import prepare  # type: ignore[import-untyped]
        except ImportError as exc:
            raise ImportError(
                "TensorFlow SavedModel export requires 'onnx' and 'onnx-tf'. "
                "Install them with:  pip install onnx onnx-tf tensorflow"
            ) from exc

        logger.info(
            "Converting ONNX (%s) to TF SavedModel (%s) ...",
            onnx_path,
            saved_model_dir,
        )

        onnx_model = onnx.load(str(onnx_path))
        tf_rep = prepare(onnx_model)
        tf_rep.export_graph(str(saved_model_dir))

        logger.info("TF SavedModel export complete: %s", saved_model_dir)
        return str(saved_model_dir)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _export_onnx(
        model: torch.nn.Module,
        sample_input: torch.Tensor,
        output_dir: Path,
        file_stem: str,
        opset: int,
        dynamic: bool,
    ) -> Path:
        """Export the PyTorch model to ONNX as an intermediate step."""
        import torch

        onnx_path = output_dir / f"{file_stem}.onnx"
        logger.info("Exporting intermediate ONNX to %s (opset %d) ...", onnx_path, opset)

        dynamic_axes = (
            {"images": {0: "batch", 2: "height", 3: "width"}}
            if dynamic
            else None
        )

        torch.onnx.export(
            model.cpu() if dynamic else model,
            sample_input.cpu() if dynamic else sample_input,
            str(onnx_path),
            verbose=False,
            opset_version=opset,
            do_constant_folding=True,
            input_names=["images"],
            output_names=["output0"],
            dynamic_axes=dynamic_axes,
        )
        return onnx_path
