"""Adapter wrapping the Ultralytics YOLO ``model.export()`` API."""

from __future__ import annotations

import logging
import os
import shutil
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class UltralyticsAdapter:
    """Wraps :meth:`ultralytics.YOLO.export` with tracking and multi-format support.

    The Ultralytics library uses its own format strings internally.
    This adapter maps canonical matrice_export format names to the Ultralytics
    equivalents and copies/moves the output artefact to the requested
    *output_dir* so that the caller has a predictable location.

    Example
    -------
    ::

        from ultralytics import YOLO
        from matrice_export.adapters.ultralytics import UltralyticsAdapter

        model = YOLO("yolov8n.pt")
        adapter = UltralyticsAdapter()
        results = adapter.export_all(model, ["onnx", "torchscript"], "./exports")
    """

    # Map from matrice_export canonical names to the string that Ultralytics
    # ``model.export(format=...)`` expects.  These are 1:1 for most
    # formats but having the explicit map keeps things future-proof and
    # lets us reject unsupported names early.
    FORMAT_MAP: dict[str, str] = {
        "torchscript": "torchscript",
        "onnx": "onnx",
        "openvino": "openvino",
        "engine": "engine",
        "coreml": "coreml",
        "saved_model": "saved_model",
        "pb": "pb",
        "tflite": "tflite",
        "edgetpu": "edgetpu",
        "tfjs": "tfjs",
        "paddle": "paddle",
        "ncnn": "ncnn",
    }

    # ------------------------------------------------------------------ #
    # Single-format export
    # ------------------------------------------------------------------ #
    def export(
        self,
        yolo_model: Any,
        fmt: str,
        device: str = "cpu",
        **kwargs: Any,
    ) -> str:
        """Call ``yolo_model.export()`` for a single format.

        Parameters
        ----------
        yolo_model:
            An Ultralytics YOLO model instance (e.g.
            ``ultralytics.YOLO("yolov8n.pt")``).
        fmt:
            Canonical format name (must be a key in :attr:`FORMAT_MAP`).
        device:
            Device string passed to Ultralytics (``"cpu"`` or ``"0"`` etc.).
        **kwargs:
            Additional keyword arguments forwarded to
            ``yolo_model.export()`` (e.g. ``half``, ``dynamic``,
            ``simplify``, ``int8``, ``nms``, ``opset``).

        Returns
        -------
        str
            Path to the exported artefact (file or directory).

        Raises
        ------
        ValueError
            If *fmt* is not a supported Ultralytics format.
        RuntimeError
            If ``yolo_model.export()`` fails.
        """
        ultralytics_fmt = self.FORMAT_MAP.get(fmt)
        if ultralytics_fmt is None:
            raise ValueError(
                f"Unsupported Ultralytics format: {fmt!r}. "
                f"Supported: {list(self.FORMAT_MAP.keys())}"
            )

        logger.info(
            "Exporting Ultralytics model to %s (device=%s, kwargs=%s)",
            ultralytics_fmt,
            device,
            kwargs,
        )

        try:
            export_path = yolo_model.export(
                format=ultralytics_fmt,
                device=device,
                **kwargs,
            )
        except Exception as exc:
            raise RuntimeError(
                f"Ultralytics export to {fmt!r} failed: {exc}"
            ) from exc

        # Ultralytics returns the path (str) to the exported artefact.
        if export_path is None:
            raise RuntimeError(
                f"Ultralytics export to {fmt!r} returned None -- "
                "check the model and format compatibility."
            )

        export_path = str(export_path)
        logger.info("Ultralytics export to %s succeeded: %s", fmt, export_path)
        return export_path

    # ------------------------------------------------------------------ #
    # Multi-format export
    # ------------------------------------------------------------------ #
    def export_all(
        self,
        yolo_model: Any,
        formats: list[str],
        output_dir: str,
        device: str = "cpu",
        **kwargs: Any,
    ) -> dict[str, dict[str, Any]]:
        """Export to multiple formats, collecting results.

        Parameters
        ----------
        yolo_model:
            An Ultralytics YOLO model instance.
        formats:
            List of canonical format names.
        output_dir:
            Directory where exported artefacts should be collected.
            If the Ultralytics export writes the artefact elsewhere it will
            be copied/moved into *output_dir*.
        device:
            Device string passed to Ultralytics.
        **kwargs:
            Forwarded to each ``export()`` call.

        Returns
        -------
        dict
            ``{format_name: {"path": str | None, "status": str, "error": str | None}}``
        """
        os.makedirs(output_dir, exist_ok=True)
        results: dict[str, dict[str, Any]] = {}

        for fmt in formats:
            try:
                raw_path = self.export(yolo_model, fmt, device=device, **kwargs)
                final_path = self._collect_artefact(raw_path, output_dir)
                results[fmt] = {
                    "path": final_path,
                    "status": "success",
                    "error": None,
                }
            except Exception as exc:
                logger.error("Ultralytics export %s failed: %s", fmt, exc, exc_info=True)
                results[fmt] = {
                    "path": None,
                    "status": "error",
                    "error": str(exc),
                }

        return results

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #
    @staticmethod
    def _collect_artefact(src_path: str, output_dir: str) -> str:
        """Copy or move the exported artefact into *output_dir*.

        If the artefact already resides inside *output_dir* it is left
        in place.

        Returns
        -------
        str
            The final path to the artefact inside *output_dir*.
        """
        src = Path(src_path).resolve()
        dest_dir = Path(output_dir).resolve()

        # Already in the target directory -- nothing to do.
        try:
            src.relative_to(dest_dir)
            return str(src)
        except ValueError:
            pass

        dest = dest_dir / src.name

        if src.is_dir():
            if dest.exists():
                shutil.rmtree(dest)
            shutil.copytree(str(src), str(dest))
        else:
            shutil.copy2(str(src), str(dest))

        logger.debug("Collected artefact %s -> %s", src, dest)
        return str(dest)
