"""PaddlePaddle (``_paddle_model``) exporter.

Converts a PyTorch model to PaddlePaddle format using ``x2paddle``.
The conversion uses ``pytorch2paddle`` in *trace* mode to translate the
PyTorch computation graph into a Paddle inference model.  A YAML metadata
sidecar is written alongside the model when ``metadata`` is provided.
"""

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 PaddlePaddleExporter(BaseExporter):
    """Export a PyTorch model to PaddlePaddle format via ``x2paddle``."""

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

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

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

        Args:
            model: PyTorch model in eval mode.
            sample_input: Sample input tensor (BCHW).
            output_dir: Directory to write the exported model directory into.
            file_stem: Base filename without extension.
            **kwargs:
                metadata (dict | None): Extra metadata to persist as a YAML
                    sidecar file inside the output directory.  Typically
                    contains ``stride`` and ``names`` keys.

        Returns:
            Absolute path to the exported PaddlePaddle model directory.
        """
        try:
            import x2paddle  # noqa: F401  # type: ignore[import-untyped]
            from x2paddle.convert import pytorch2paddle  # type: ignore[import-untyped]
        except ImportError as exc:
            raise ImportError(
                "PaddlePaddle export requires 'paddlepaddle' and 'x2paddle'. "
                "Install them with:  pip install paddlepaddle x2paddle"
            ) from exc

        metadata: dict | None = kwargs.get("metadata", None)

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

        logger.info(
            "PaddlePaddle export: converting with x2paddle %s ...",
            getattr(x2paddle, "__version__", "unknown"),
        )

        pytorch2paddle(
            module=model,
            save_dir=str(paddle_dir),
            jit_type="trace",
            input_examples=[sample_input],
        )

        # ------------------------------------------------------------------
        # Optional metadata sidecar (YAML)
        # ------------------------------------------------------------------
        if metadata is not None:
            try:
                import yaml
            except ImportError:
                logger.warning(
                    "PaddlePaddle export: PyYAML not installed -- skipping "
                    "metadata YAML sidecar."
                )
            else:
                yaml_path = paddle_dir / f"{file_stem}.yaml"
                with open(yaml_path, "w", encoding="utf-8") as fh:
                    yaml.safe_dump(metadata, fh, default_flow_style=False)
                logger.info("PaddlePaddle export: metadata saved to %s", yaml_path)

        logger.info("PaddlePaddle export: saved to %s", paddle_dir)
        return str(paddle_dir)
