"""TorchScript model exporter."""

from __future__ import annotations

import json
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 TorchScriptExporter(BaseExporter):
    """Export a PyTorch model to TorchScript format.

    Traces the model with ``torch.jit.trace`` and persists metadata
    (input shape, stride, class names) in the archive's ``extra_files``.
    Optionally applies ``optimize_for_mobile`` for lite-interpreter usage.
    """

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

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

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

        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:
                optimize (bool): If ``True``, apply ``optimize_for_mobile``
                    and save for the lite interpreter.  Defaults to ``False``.
                stride (int | None): Model stride value to embed in metadata.
                names (dict | list | None): Class names to embed in metadata.

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

        optimize: bool = kwargs.get("optimize", False)
        stride: int | None = kwargs.get("stride", None)
        names = kwargs.get("names", None)

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

        # Build metadata payload
        metadata: dict = {
            "shape": list(sample_input.shape),
            "stride": stride if stride is not None else int(max(model.stride)) if hasattr(model, "stride") else 1,
            "names": names,
        }
        extra_files: dict[str, str] = {"config.txt": json.dumps(metadata)}

        logger.info(
            "TorchScript export: tracing model with torch %s ...",
            torch.__version__,
        )

        ts = torch.jit.trace(model, sample_input, strict=False)

        if optimize:
            try:
                from torch.utils.mobile_optimizer import optimize_for_mobile
            except ImportError as exc:
                raise ImportError(
                    "torch.utils.mobile_optimizer is required for "
                    "optimize_for_mobile but could not be imported. "
                    "Ensure your PyTorch build includes mobile optimisation support."
                ) from exc

            logger.info("TorchScript export: applying optimize_for_mobile ...")
            optimized = optimize_for_mobile(ts)
            optimized._save_for_lite_interpreter(
                str(output_path), _extra_files=extra_files
            )
        else:
            ts.save(str(output_path), _extra_files=extra_files)

        logger.info("TorchScript export: saved %s", output_path)
        return str(output_path)
