"""TensorFlow GraphDef (frozen ``.pb``) exporter.

Freezes a TF SavedModel into a single ``.pb`` file by converting all
variables to constants (``convert_variables_to_constants_v2``).
"""

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 TFGraphDefExporter(BaseExporter):
    """Export a TF SavedModel to a frozen GraphDef ``.pb`` file."""

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

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

    # ------------------------------------------------------------------
    # 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 to TF GraphDef (``.pb``).

        This exporter depends on a TF SavedModel having already been created.
        Supply the path via the ``saved_model_path`` keyword argument, or let
        the exporter create one automatically by chaining to
        :class:`TFSavedModelExporter`.

        Keyword Args:
            saved_model_path (str | Path | None): Path to an existing TF
                SavedModel directory.  When *None*, a SavedModel is produced
                automatically via :class:`TFSavedModelExporter`.
            onnx_path (str | Path | None): Forwarded to TFSavedModelExporter
                when ``saved_model_path`` is not supplied.

        Returns:
            Path to the exported ``.pb`` file.
        """
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        pb_path = output_dir / f"{file_stem}{self.suffix}"

        saved_model_path: str | Path | None = kwargs.get("saved_model_path")

        # If no saved-model directory was provided, build one first.
        if saved_model_path is None:
            from matrice_export.formats.tf_saved import TFSavedModelExporter

            logger.info("No saved_model_path supplied — creating TF SavedModel first ...")
            saved_model_path = TFSavedModelExporter().export(
                model, sample_input, output_dir, file_stem, **kwargs
            )

        saved_model_path = Path(saved_model_path)
        if not saved_model_path.exists():
            raise FileNotFoundError(
                f"TF SavedModel directory not found: {saved_model_path}"
            )

        # --- freeze the saved model -----------------------------------
        try:
            import tensorflow as tf
            from tensorflow.python.framework.convert_to_constants import (
                convert_variables_to_constants_v2,
            )
        except ImportError as exc:
            raise ImportError(
                "GraphDef export requires TensorFlow. "
                "Install it with:  pip install tensorflow"
            ) from exc

        logger.info(
            "Freezing TF SavedModel (%s) to GraphDef (%s) ...",
            saved_model_path,
            pb_path,
        )

        # Load the saved model and obtain a concrete function.
        loaded = tf.saved_model.load(str(saved_model_path))
        concrete_func = loaded.signatures[
            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
        ]
        frozen_func = convert_variables_to_constants_v2(concrete_func)
        frozen_func.graph.as_graph_def()

        tf.io.write_graph(
            graph_or_graph_def=frozen_func.graph,
            logdir=str(pb_path.parent),
            name=pb_path.name,
            as_text=False,
        )

        logger.info("GraphDef export complete: %s", pb_path)
        return str(pb_path)
