"""Unified export pipeline supporting PyTorch, ONNX, and Ultralytics models."""

from __future__ import annotations

import logging
import os
from collections.abc import Callable
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import numpy as np

    from matrice_export.formats.base import BaseExporter

logger = logging.getLogger(__name__)

_REGISTERED = False


def _is_pytorch_model(model: Any) -> bool:
    """Check if model is a PyTorch nn.Module without hard-importing torch."""
    try:
        import torch
        return isinstance(model, torch.nn.Module)
    except ImportError:
        return False


def _is_ultralytics_model(model: Any) -> bool:
    """Check if model looks like an Ultralytics YOLO object.

    Ultralytics models expose ``.export()`` and ``.predict()`` methods.
    We avoid importing ultralytics directly so it stays an optional dependency.
    """
    return hasattr(model, "export") and hasattr(model, "predict") and not _is_pytorch_model(model)


def _is_onnx_path(model: Any) -> bool:
    """Check if model is a path (str or Path) pointing to an ``.onnx`` file."""
    if isinstance(model, (str, Path)):
        return str(model).endswith(".onnx")
    return False


class ExportPipeline:
    """Unified export pipeline supporting PyTorch, ONNX, and Ultralytics models.

    Dispatches export work based on the detected model type:

    * **PyTorch** (``torch.nn.Module``) -- uses registered :class:`BaseExporter`
      format classes from the ``EXPORTERS`` registry.
    * **ONNX** (``str`` / ``Path`` ending in ``.onnx``) -- delegates to
      :class:`~matrice_export.adapters.onnx_input.OnnxInputAdapter`.
    * **Ultralytics YOLO** (object with ``.export`` and ``.predict``) -- delegates
      to :class:`~matrice_export.adapters.ultralytics.UltralyticsAdapter`.

    Parameters
    ----------
    model:
        A PyTorch ``nn.Module``, a path to an ``.onnx`` file, or an
        Ultralytics YOLO model object.
    sample_input:
        Optional sample input tensor (``torch.Tensor`` or ``numpy.ndarray``).
        Used to generate a PyTorch baseline output for post-export validation.
    task:
        The model task type, e.g. ``"detection"``, ``"classification"``,
        ``"segmentation"``.  Informational; passed to adapters/validators.
    tracker:
        Optional :class:`~matrice_export.tracking.ExportTracker` instance for
        reporting status to the Matrice backend.
    """

    # ------------------------------------------------------------------ #
    # Class-level registries
    # ------------------------------------------------------------------ #
    EXPORTERS: dict[str, type[BaseExporter]] = {}
    VALIDATORS: dict[str, Callable] = {}

    # ------------------------------------------------------------------ #
    # Construction helpers
    # ------------------------------------------------------------------ #
    def __init__(
        self,
        model: Any,
        sample_input: Any | None = None,
        task: str = "detection",
        tracker: Any | None = None,
    ) -> None:
        self._auto_register()
        self.model = model
        self.task = task
        self.tracker = tracker
        self.sample_input = sample_input
        self.baseline_output: np.ndarray | None = None

        # Detect model type ------------------------------------------------
        if _is_pytorch_model(model):
            self.model_type = "pytorch"
        elif _is_ultralytics_model(model):
            self.model_type = "ultralytics"
        elif _is_onnx_path(model):
            self.model_type = "onnx"
            if not os.path.isfile(str(model)):
                raise FileNotFoundError(f"ONNX model not found: {model}")
        else:
            raise TypeError(
                f"Unsupported model type: {type(model).__name__}. "
                "Expected a torch.nn.Module, an Ultralytics YOLO model, "
                "or a path to an .onnx file."
            )

        # Generate baseline output for later validation --------------------
        if sample_input is not None and self.model_type == "pytorch":
            self.baseline_output = self._generate_baseline(model, sample_input)

    # ------------------------------------------------------------------ #
    # Registry helpers
    # ------------------------------------------------------------------ #
    @classmethod
    def _auto_register(cls) -> None:
        """Auto-register all exporters and validators on first use."""
        global _REGISTERED
        if _REGISTERED:
            return
        _REGISTERED = True

        # -- Exporters -------------------------------------------------------
        from matrice_export.formats import get_all_exporters

        for exporter_class in get_all_exporters():
            cls.register_exporter(exporter_class)

        # -- Validators -------------------------------------------------------
        from matrice_export.validators import get_all_validators

        for fmt_name, validator_class in get_all_validators().items():
            cls.register_validator(fmt_name, validator_class().validate)

    @classmethod
    def register_exporter(cls, exporter_class: type[BaseExporter]) -> type[BaseExporter]:
        """Register a :class:`BaseExporter` subclass in the format registry.

        Can be used as a decorator::

            @ExportPipeline.register_exporter
            class MyExporter(BaseExporter):
                ...

        Or called directly::

            ExportPipeline.register_exporter(MyExporter)
        """
        instance = exporter_class()
        name = instance.format_name
        cls.EXPORTERS[name] = exporter_class
        logger.debug("Registered exporter %r for format %r", exporter_class.__name__, name)
        return exporter_class

    @classmethod
    def register_validator(cls, format_name: str, validator_fn: Callable) -> None:
        """Register a validation function for *format_name*.

        Parameters
        ----------
        format_name:
            The canonical format name (e.g. ``"onnx"``).
        validator_fn:
            A callable with signature
            ``(model_path: str, sample_input: np.ndarray, baseline: np.ndarray, **kw) -> dict``
            returning a validation result dictionary.
        """
        cls.VALIDATORS[format_name] = validator_fn
        logger.debug("Registered validator for format %r", format_name)

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #
    def export(
        self,
        formats: list[str],
        output_dir: str,
        validate: bool = True,
        **kwargs: Any,
    ) -> dict[str, dict[str, Any]]:
        """Export the model to the requested *formats*.

        Parameters
        ----------
        formats:
            List of format names (e.g. ``["onnx", "torchscript", "openvino"]``).
        output_dir:
            Directory where exported artefacts will be written.
        validate:
            If ``True``, run the registered validator (if any) after each
            successful export.
        **kwargs:
            Extra keyword arguments forwarded to the underlying exporter (e.g.
            ``half=True``, ``dynamic=True``, ``opset=13``).

        Returns
        -------
        dict
            ``{format_name: {"path": str, "status": str, "validation": dict | None}}``
            One entry per requested format.  ``status`` is ``"success"`` or
            ``"error"``.  On error the dict also contains an ``"error"`` key.
        """
        os.makedirs(output_dir, exist_ok=True)

        if self.tracker is not None:
            self.tracker.report_pipeline_start(formats)

        dispatch = {
            "pytorch": self._export_pytorch,
            "ultralytics": self._export_ultralytics,
            "onnx": self._export_from_onnx,
        }
        handler = dispatch[self.model_type]
        results = handler(formats, output_dir, validate, **kwargs)

        if self.tracker is not None:
            self.tracker.report_pipeline_complete(results)

        return results

    # ------------------------------------------------------------------ #
    # Private dispatch methods
    # ------------------------------------------------------------------ #
    def _export_pytorch(
        self,
        formats: list[str],
        output_dir: str,
        validate: bool,
        **kwargs: Any,
    ) -> dict[str, dict[str, Any]]:
        """Export a PyTorch ``nn.Module`` via registered :class:`BaseExporter` classes."""
        import numpy as np
        import torch

        results: dict[str, dict[str, Any]] = {}

        for fmt in formats:
            if self.tracker is not None:
                self.tracker.report_export_start(fmt)
            try:
                exporter_cls = self.EXPORTERS.get(fmt)
                if exporter_cls is None:
                    raise ValueError(
                        f"No exporter registered for format {fmt!r}. "
                        f"Available: {list(self.EXPORTERS.keys())}"
                    )
                exporter = exporter_cls()

                # Ensure a sample input is available
                sample = self.sample_input
                if sample is None:
                    raise ValueError(
                        f"sample_input is required to export PyTorch model to {fmt!r}"
                    )

                # Normalise sample to torch.Tensor
                if isinstance(sample, np.ndarray):
                    sample = torch.from_numpy(sample)

                model_path = exporter.export(
                    model=self.model,
                    sample_input=sample,
                    output_dir=output_dir,
                    **kwargs,
                )

                # Validation -------------------------------------------
                validation = None
                if validate:
                    validation = self._validate(fmt, model_path)

                results[fmt] = {
                    "path": str(model_path),
                    "status": "success",
                    "validation": validation,
                }

                if self.tracker is not None:
                    self.tracker.report_export_success(fmt, str(model_path), validation)

                logger.info("Successfully exported %s to %s", fmt, model_path)

            except Exception as exc:
                logger.error("Export to %s failed: %s", fmt, exc, exc_info=True)
                results[fmt] = {
                    "path": None,
                    "status": "error",
                    "error": str(exc),
                    "validation": None,
                }
                if self.tracker is not None:
                    self.tracker.report_export_failure(fmt, str(exc))

        return results

    def _export_ultralytics(
        self,
        formats: list[str],
        output_dir: str,
        validate: bool,
        **kwargs: Any,
    ) -> dict[str, dict[str, Any]]:
        """Delegate to :class:`UltralyticsAdapter` then run validation on top."""
        from matrice_export.adapters.ultralytics import UltralyticsAdapter

        adapter = UltralyticsAdapter()
        device = kwargs.pop("device", "cpu")

        raw_results = adapter.export_all(
            yolo_model=self.model,
            formats=formats,
            output_dir=output_dir,
            device=device,
            **kwargs,
        )

        results: dict[str, dict[str, Any]] = {}

        for fmt, entry in raw_results.items():
            if self.tracker is not None and entry["status"] == "success":
                self.tracker.report_export_success(fmt, entry["path"])
            elif self.tracker is not None:
                self.tracker.report_export_failure(fmt, entry.get("error", "unknown"))

            validation = None
            if validate and entry["status"] == "success":
                validation = self._validate(fmt, entry["path"])

            results[fmt] = {
                "path": entry.get("path"),
                "status": entry["status"],
                "error": entry.get("error"),
                "validation": validation,
            }

        return results

    def _export_from_onnx(
        self,
        formats: list[str],
        output_dir: str,
        validate: bool,
        **kwargs: Any,
    ) -> dict[str, dict[str, Any]]:
        """Delegate to :class:`OnnxInputAdapter` for ONNX -> downstream conversions."""
        from matrice_export.adapters.onnx_input import OnnxInputAdapter

        adapter = OnnxInputAdapter()
        onnx_path = str(self.model)
        results: dict[str, dict[str, Any]] = {}

        for fmt in formats:
            if self.tracker is not None:
                self.tracker.report_export_start(fmt)
            try:
                model_path = adapter.export(
                    onnx_path=onnx_path,
                    fmt=fmt,
                    output_dir=output_dir,
                    **kwargs,
                )
                validation = None
                if validate:
                    validation = self._validate(fmt, model_path)

                results[fmt] = {
                    "path": str(model_path),
                    "status": "success",
                    "validation": validation,
                }

                if self.tracker is not None:
                    self.tracker.report_export_success(fmt, str(model_path), validation)

                logger.info("Successfully converted ONNX -> %s at %s", fmt, model_path)

            except Exception as exc:
                logger.error("ONNX -> %s conversion failed: %s", fmt, exc, exc_info=True)
                results[fmt] = {
                    "path": None,
                    "status": "error",
                    "error": str(exc),
                    "validation": None,
                }
                if self.tracker is not None:
                    self.tracker.report_export_failure(fmt, str(exc))

        return results

    # ------------------------------------------------------------------ #
    # Validation helper
    # ------------------------------------------------------------------ #
    def _validate(
        self,
        format_name: str,
        model_path: str,
    ) -> dict | None:
        """Run the registered validator for *format_name* if one exists.

        Returns the validator result dict, or ``None`` if no validator is
        registered or if the required inputs (``sample_input`` /
        ``baseline_output``) are not available.
        """
        import numpy as np

        validator_fn = self.VALIDATORS.get(format_name)
        if validator_fn is None:
            logger.debug("No validator registered for %r, skipping.", format_name)
            return None

        if self.sample_input is None:
            logger.debug("No sample_input; skipping validation for %r.", format_name)
            return None

        # Ensure numpy arrays for validator
        sample_np = self.sample_input
        if not isinstance(sample_np, np.ndarray):
            # Handle torch.Tensor or similar
            if hasattr(sample_np, "detach"):
                sample_np = sample_np.detach().cpu().numpy()
            else:
                sample_np = np.asarray(sample_np)

        baseline_np = self.baseline_output  # may be None

        try:
            result = validator_fn(
                model_path=str(model_path),
                sample_input=sample_np,
                baseline=baseline_np,
                task=self.task,
            )
            logger.info("Validation for %r: %s", format_name, result)
            return result
        except Exception as exc:
            logger.warning("Validator for %r raised an exception: %s", format_name, exc, exc_info=True)
            return {"status": "error", "error": str(exc)}

    # ------------------------------------------------------------------ #
    # Static helpers
    # ------------------------------------------------------------------ #
    @staticmethod
    def available_formats() -> list[dict[str, Any]]:
        """Return metadata for every format that has a registered exporter.

        Each element is a dict with keys:

        * ``name`` -- canonical format name
        * ``suffix`` -- file extension or directory suffix
        * ``requires_gpu`` -- whether a GPU is needed for export
        * ``has_validator`` -- whether a validator is registered
        """
        from matrice_export.adapters.onnx_input import OnnxInputAdapter

        formats: list[dict[str, Any]] = []

        # Formats from BaseExporter registry
        for name, exporter_cls in ExportPipeline.EXPORTERS.items():
            inst = exporter_cls()
            formats.append(
                {
                    "name": inst.format_name,
                    "suffix": inst.suffix,
                    "requires_gpu": inst.requires_gpu,
                    "has_validator": name in ExportPipeline.VALIDATORS,
                }
            )

        # Formats reachable via OnnxInputAdapter (ONNX -> downstream)
        existing_names = {f["name"] for f in formats}
        for onnx_fmt in OnnxInputAdapter.SUPPORTED_FORMATS:
            if onnx_fmt not in existing_names:
                formats.append(
                    {
                        "name": onnx_fmt,
                        "suffix": None,
                        "requires_gpu": onnx_fmt == "engine",
                        "has_validator": onnx_fmt in ExportPipeline.VALIDATORS,
                        "source": "onnx_adapter",
                    }
                )

        return formats

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #
    @staticmethod
    def _generate_baseline(model: Any, sample_input: Any) -> np.ndarray | None:
        """Run forward pass on the PyTorch model to capture baseline output.

        Returns
        -------
        numpy.ndarray or None
            The model output converted to a numpy array, or ``None`` if
            the forward pass fails.
        """
        try:
            import numpy as np
            import torch

            was_training = model.training
            model.eval()

            if isinstance(sample_input, np.ndarray):
                tensor_input = torch.from_numpy(sample_input)
            else:
                tensor_input = sample_input

            with torch.no_grad():
                output = model(tensor_input)

            if was_training:
                model.train()

            # Handle tuple / list outputs (take first element)
            if isinstance(output, (tuple, list)):
                output = output[0]

            # Handle objects with .logits (e.g. HuggingFace model outputs)
            if hasattr(output, "logits"):
                output = output.logits

            return output.detach().cpu().numpy()
        except Exception as exc:
            logger.warning("Could not generate baseline output: %s", exc)
            return None
