"""TensorFlow.js (``_web_model``) exporter.

Converts a frozen TF GraphDef (``.pb``) to a TensorFlow.js web model using
the ``tensorflowjs_converter`` CLI.
"""

from __future__ import annotations

import logging
import re
import subprocess
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 TFJSExporter(BaseExporter):
    """Export a model to TensorFlow.js web model format."""

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

    @property
    def suffix(self) -> str:
        return "_web_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 to TensorFlow.js.

        This exporter requires a frozen GraphDef (``.pb``) file as input.
        Supply it via the ``pb_path`` keyword argument, or let the exporter
        create one by chaining through the TF export pipeline.

        Keyword Args:
            pb_path (str | Path | None): Path to an existing ``.pb`` file.
                When *None*, one is produced via :class:`TFGraphDefExporter`.
            int8 (bool): Apply uint8 quantisation in the JS converter.
                Defaults to ``False``.
            Any additional kwargs are forwarded to upstream exporters if
            the ``.pb`` file needs to be created.

        Returns:
            Path to the exported ``*_web_model`` directory.
        """
        try:
            import tensorflowjs  # noqa: F401  # type: ignore[import-untyped]
        except ImportError as exc:
            raise ImportError(
                "TensorFlow.js export requires 'tensorflowjs'. "
                "Install it with:  pip install tensorflowjs"
            ) from exc

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

        pb_path: str | Path | None = kwargs.get("pb_path")
        int8: bool = kwargs.get("int8", False)

        # Ensure a frozen GraphDef is available.
        if pb_path is None:
            from matrice_export.formats.tf_graphdef import TFGraphDefExporter

            logger.info(
                "No pb_path supplied — creating TF GraphDef first ..."
            )
            pb_path = TFGraphDefExporter().export(
                model, sample_input, output_dir, file_stem, **kwargs
            )

        pb_path = Path(pb_path)
        if not pb_path.exists():
            raise FileNotFoundError(f"GraphDef .pb file not found: {pb_path}")

        logger.info(
            "Converting GraphDef (%s) to TF.js (%s) ...", pb_path, tfjs_dir
        )

        # Build tensorflowjs_converter CLI arguments.
        args = [
            "tensorflowjs_converter",
            "--input_format=tf_frozen_model",
        ]
        if int8:
            args.append("--quantize_uint8")
        args += [
            "--output_node_names=Identity,Identity_1,Identity_2,Identity_3",
            str(pb_path),
            str(tfjs_dir),
        ]

        subprocess.run(args, check=True)

        # Sort Identity output nodes in the model.json so they appear in
        # ascending order (Identity, Identity_1, Identity_2, Identity_3).
        self._sort_identity_outputs(tfjs_dir / "model.json")

        logger.info("TF.js export complete: %s", tfjs_dir)
        return str(tfjs_dir)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _sort_identity_outputs(json_path: Path) -> None:
        """Re-order ``Identity*`` keys inside *model.json* to ascending order.

        The ``tensorflowjs_converter`` sometimes emits the Identity output
        entries in a non-deterministic order.  This helper normalises them so
        downstream consumers can rely on a stable ordering.
        """
        if not json_path.exists():
            return

        text = json_path.read_text()
        substituted = re.sub(
            r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
            r'"Identity.?.?": {"name": "Identity.?.?"}, '
            r'"Identity.?.?": {"name": "Identity.?.?"}, '
            r'"Identity.?.?": {"name": "Identity.?.?"}}}',
            r'{"outputs": {"Identity": {"name": "Identity"}, '
            r'"Identity_1": {"name": "Identity_1"}, '
            r'"Identity_2": {"name": "Identity_2"}, '
            r'"Identity_3": {"name": "Identity_3"}}}',
            text,
        )
        json_path.write_text(substituted)
