"""Base Trainer class — SageMaker Estimator-like orchestrator for BYOM training jobs.

Provides ``Trainer`` (the abstract base), ``TrainResult`` (the output
data model), and ``TrainerError`` (lifecycle exception).  Every BYOM
repository uses one of the two concrete sub-classes instead of this
class directly:

* ``LoopTrainer``      — you own the PyTorch training for-loop.
* ``DelegatedTrainer`` — the framework owns the loop (Ultralytics, D-FINE …).
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import TYPE_CHECKING, Any

from matrice_models.common.model_loading import load_model_with_strategies as common_load_model_with_strategies
from pydantic import BaseModel, Field


if TYPE_CHECKING:
    from collections.abc import Callable

    import torch

    from matrice_models.common.model_loading import ModelLoadStrategies
    from matrice_models.config.base import ModelInfo, TrainConfig
    from matrice_models.training.callbacks import TrainerCallback


logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
#  Data models                                                         #
# ------------------------------------------------------------------ #


class TrainResult(BaseModel):
    """Output of a completed training run.

    Captures the best metric achieved, the number of epochs that ran,
    saved checkpoint paths, and the evaluation payload that was (or will
    be) sent to ActionTracker.
    """

    best_metric_name: str = Field(
        description="Name of the tracked metric (e.g. acc@1, mAP50)",
    )
    best_metric_value: float = Field(
        description="Best value achieved for the tracked metric",
    )
    best_epoch: int = Field(
        ge=0,
        description="Epoch at which the best value was recorded",
    )
    epochs_completed: int = Field(
        ge=0,
        description="Total number of epochs that ran before stopping",
    )
    checkpoint_paths: list[Path] = Field(
        default_factory=list,
        description="Paths to saved checkpoint files",
    )
    evaluation_payload: list[dict[str, Any]] = Field(
        default_factory=list,
        description="Evaluation results in ActionTracker format",
    )


# ------------------------------------------------------------------ #
#  Exceptions                                                          #
# ------------------------------------------------------------------ #


class TrainerError(Exception):
    """Raised when a training lifecycle phase fails.

    Attributes:
        phase: The ActionTracker status code of the failed phase
            (e.g. ``"MDL_TRN_DTL"``).
    """

    def __init__(self, phase: str, message: str) -> None:
        """Initialize with the failing phase and error message.

        Args:
            phase: ActionTracker status code (e.g. ``MDL_TRN_DTL``).
            message: Human-readable description of the failure.
        """
        self.phase = phase
        super().__init__(f"[{phase}] {message}")


# ------------------------------------------------------------------ #
#  Trainer base class                                                  #
# ------------------------------------------------------------------ #


class Trainer(ABC):
    """Abstract base for all BYOM training jobs.

    Modelled after AWS SageMaker's ``Estimator``: call :meth:`fit` and
    the full lifecycle — data loading, model initialisation, training,
    checkpointing, evaluation — runs with ActionTracker status updates
    at every phase.

    Subclass one of the concrete variants instead of this class directly:

    * ``LoopTrainer`` — you own the PyTorch for-loop
      (classification, R2Plus1D).
    * ``DelegatedTrainer`` — the framework owns the loop
      (Ultralytics YOLO, D-FINE solver).

    Each subclass only implements a handful of hooks; all boilerplate
    (ActionTracker wiring, error handling, checkpoint upload) is handled
    here.

    Example::

        class MyTrainer(LoopTrainer):
            def setup_data(self) -> None: ...
            def setup_model(self) -> None: ...

            ...


        trainer = MyTrainer(action_id, model_info, config)
        result = trainer.fit()
    """

    # ------------------------------------------------------------------ #
    #  Construction                                                        #
    # ------------------------------------------------------------------ #

    def __init__(
        self,
        action_id: str | None,
        model_info: ModelInfo,
        config: TrainConfig,
        *,
        checkpoint_dir: str | Path = "checkpoints",
        callbacks: list[TrainerCallback] | None = None,
    ) -> None:
        """Initialize the trainer.

        Args:
            action_id: Platform action ID.  Pass ``None`` for local /
                offline runs (a ``LocalActionTracker`` or no-op stub
                will be created automatically).
            model_info: Model identity and capabilities.
            config: Training hyperparameters.
            checkpoint_dir: Directory where checkpoints are written.
            callbacks: Optional list of :class:`TrainerCallback`
                instances.  When ``None`` the default set
                (``ActionTrackerCallback``) is registered automatically.
        """
        self.model_info = model_info
        self.config = config
        self.checkpoint_dir = Path(checkpoint_dir)
        self.callbacks: list[TrainerCallback] = callbacks if callbacks is not None else self._default_callbacks()

        self.model: Any = None
        self.device: torch.device | None = None

        self._result: TrainResult | None = None
        self._label_map: dict[str, str] = {}

        self.action_tracker: Any = _create_action_tracker(
            action_id,
            architecture=model_info.architecture,
            task=model_info.task,
        )

    # ------------------------------------------------------------------ #
    #  Public entry-point                                                  #
    # ------------------------------------------------------------------ #

    def fit(self) -> TrainResult:
        """Run the full training lifecycle.

        This is the single entry-point, analogous to
        ``sagemaker.estimator.Estimator.fit()``.  The phases are:

        1. **Acknowledge** — ``MDL_TRN_ACK``
        2. **Data loading** — ``MDL_TRN_DTL``
        3. **Model init** — ``MDL_TRN_MDL``
        4. **Training loop** — ``MDL_TRN_CMPL``
        5. **Checkpoint upload** — ``MDL_TRN_BMS``
        6. **Evaluation** — ``MDL_TRN_EVL``

        Each phase is wrapped in ActionTracker status updates.  On
        failure the tracker receives an ``ERROR`` status and a
        :class:`TrainerError` is raised.

        Returns:
            A :class:`TrainResult` summarising the completed run.

        Raises:
            TrainerError: If any lifecycle phase fails.
        """
        self._run_phase("MDL_TRN_ACK", "Model training acknowledged")

        self._run_phase(
            "MDL_TRN_DTL",
            "Training dataset loaded",
            self._do_setup_data,
        )
        self._run_phase(
            "MDL_TRN_MDL",
            "Model loaded",
            self._do_setup_model,
        )
        self._run_phase(
            "MDL_TRN_CMPL",
            "Model training completed",
            self.train_loop,
        )
        self._run_phase(
            "MDL_TRN_BMS",
            "Best model saved",
            self._do_save_model,
        )
        self._run_phase("MDL_TRN_SUCCESS", "Model training successful")
        self._run_phase(
            "MDL_TRN_EVL",
            "Model training and evaluation complete",
            self._do_evaluate,
            success_level="SUCCESS",
        )

        assert self._result is not None  # populated by train_loop()
        return self._result

    # ------------------------------------------------------------------ #
    #  Abstract hooks — subclasses MUST implement                          #
    # ------------------------------------------------------------------ #

    @abstractmethod
    def setup_data(self) -> None:
        """Load datasets and create data loaders.

        Implementations should store loaders as instance attributes
        (e.g. ``self.train_loader``) so that :meth:`train_loop` and
        :meth:`run_evaluation` can access them.
        """

    @abstractmethod
    def setup_model(self) -> None:
        """Construct or load the model and assign it to ``self.model``.

        The compute device is resolved automatically *after* this hook
        returns.  Moving the model to the device is the responsibility
        of the subclass (in this hook or in :meth:`train_loop`), because
        different frameworks handle device placement differently.

        For checkpoint/predefined selection flows, subclasses can use
        :meth:`load_model_with_strategies`.
        """

    @abstractmethod
    def get_label_map(self) -> dict[str, str]:
        """Return an index-to-label mapping for the ActionTracker.

        Returns:
            Dict mapping **string** class indices to human-readable
            class names, e.g. ``{"0": "cat", "1": "dog"}``.
        """

    @abstractmethod
    def train_loop(self) -> None:
        """Execute the main training loop.

        ``LoopTrainer`` provides a standard PyTorch epoch-loop.
        ``DelegatedTrainer`` delegates to the framework's own API.

        **Contract:** implementations *must* populate ``self._result``
        (a :class:`TrainResult`) before returning.
        """

    @abstractmethod
    def run_evaluation(self) -> list[dict[str, Any]]:
        """Run post-training evaluation and return the results payload.

        The returned list is forwarded to
        ``action_tracker.save_evaluation_results()``.

        Returns:
            Metric dicts in ActionTracker format::

                [{"splitType": "val", "metricName": "acc@1", "metricValue": 0.95, "category": "all"}, ...]
        """

    def get_model_load_strategies(self) -> ModelLoadStrategies | None:
        """Optionally provide shared model-loading strategies.

        Returns:
            ``ModelLoadStrategies`` to enable automatic model-loading in
            :meth:`_do_setup_model`, or ``None`` to use subclass
            :meth:`setup_model` behavior.
        """
        return None

    # ------------------------------------------------------------------ #
    #  Overridable hooks — subclasses MAY override                         #
    # ------------------------------------------------------------------ #

    def save_model(self) -> None:
        """Save the best model and upload checkpoints to the platform.

        The default implementation uploads ``model_best.pth.tar`` and
        ``model_best.pt`` from :attr:`checkpoint_dir`.  If the ``.pt``
        file does not exist but ``self.model`` is set, it is serialised
        with ``torch.save`` first.

        Override for frameworks that use a different checkpoint layout
        (e.g. Ultralytics writes to ``runs/detect/train/weights/``).
        """
        import torch

        best_pth = self.checkpoint_dir / "model_best.pth.tar"
        best_pt = self.checkpoint_dir / "model_best.pt"

        if best_pth.exists():
            self.action_tracker.upload_checkpoint(str(best_pth))

        if best_pt.exists():
            self.action_tracker.upload_checkpoint(str(best_pt))
        elif self.model is not None:
            self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
            torch.save(self.model, str(best_pt))
            self.action_tracker.upload_checkpoint(str(best_pt))

    def load_model_with_strategies(self, strategies: ModelLoadStrategies) -> Any:
        """Load model using ActionTracker checkpoint/predefined strategy.

        Args:
            strategies: Strategy callbacks for checkpoint/predefined/scratch branches.

        Returns:
            Loaded model artifact selected from the first available branch.
        """
        return common_load_model_with_strategies(self.action_tracker, strategies)

    def log_epoch(self, epoch: int, metrics: list[dict[str, Any]]) -> None:
        """Forward per-epoch metrics to the ActionTracker.

        Called by ``LoopTrainer`` (or your own loop) at the end of each
        epoch.  Override to add side-channels such as TensorBoard or
        Weights & Biases logging.

        Args:
            epoch: Zero-based epoch number.
            metrics: List of metric dicts, each containing at minimum
                ``splitType``, ``metricName``, and ``metricValue``.
        """
        self.action_tracker.log_epoch_results(epoch, metrics)

    # ------------------------------------------------------------------ #
    #  Callback helpers                                                    #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _default_callbacks() -> list[TrainerCallback]:
        """Return the callbacks registered when the caller passes ``None``.

        Returns:
            A list containing a single :class:`ActionTrackerCallback`.
        """
        from matrice_models.training.callbacks import ActionTrackerCallback

        return [ActionTrackerCallback()]

    def notify(self, hook: str, *args: Any, **kwargs: Any) -> None:
        """Fire a named callback hook on every registered callback.

        Exceptions from individual callbacks are logged but do **not**
        interrupt the training loop.

        Args:
            hook: Method name to invoke (e.g. ``"on_epoch_end"``).
            *args: Positional arguments forwarded to the hook.
            **kwargs: Keyword arguments forwarded to the hook.
        """
        from matrice_models.training.callbacks import run_callbacks

        run_callbacks(self.callbacks, hook, *args, **kwargs)

    # ------------------------------------------------------------------ #
    #  Internal lifecycle wiring                                           #
    # ------------------------------------------------------------------ #

    def _run_phase(
        self,
        status_code: str,
        ok_message: str,
        fn: Callable[[], None] | None = None,
        *,
        success_level: str = "OK",
    ) -> None:
        """Execute *fn* inside an ActionTracker-guarded lifecycle phase.

        On success the tracker receives
        ``(status_code, success_level, ok_message)``.
        On failure it receives ``(status_code, 'ERROR', str(exc))`` and
        a :class:`TrainerError` is raised so ``fit()`` aborts.

        Args:
            status_code: ActionTracker status code for this phase.
            ok_message: Human-readable message sent on success.
            fn: Work to execute.  ``None`` means status-only update.
            success_level: Level string on success (default ``"OK"``).
        """
        try:
            if fn is not None:
                fn()
            self.action_tracker.update_status(
                status_code,
                success_level,
                ok_message,
            )
        except TrainerError:
            raise
        except Exception as exc:
            logger.exception("Phase %s failed", status_code)
            self.action_tracker.update_status(
                status_code,
                "ERROR",
                str(exc),
            )
            self.action_tracker.log_error(
                __file__,
                f"Trainer._run_phase/{status_code}",
                str(exc),
            )
            raise TrainerError(status_code, str(exc)) from exc

    # -- composite phase helpers --

    def _do_setup_data(self) -> None:
        """Run ``setup_data`` then register the label map."""
        self.setup_data()
        self._label_map = self.get_label_map()
        self.action_tracker.add_index_to_category(self._label_map)

    def _do_setup_model(self) -> None:
        """Run ``setup_model`` then resolve the compute device."""
        strategies = self.get_model_load_strategies()
        if strategies is not None:
            self.model = self.load_model_with_strategies(strategies)
        else:
            self.setup_model()
        from matrice_models.training.device import resolve_device

        self.device = resolve_device(self.config.device)

    def _do_save_model(self) -> None:
        """Delegate to the (possibly overridden) ``save_model``."""
        self.save_model()

    def _do_evaluate(self) -> None:
        """Run evaluation and persist the results via ActionTracker."""
        payload = self.run_evaluation()
        self.action_tracker.save_evaluation_results(payload)
        if self._result is not None:
            self._result.evaluation_payload = payload


# ------------------------------------------------------------------ #
#  Module-level helpers                                                #
# ------------------------------------------------------------------ #


def _create_action_tracker(
    action_id: str | None,
    *,
    architecture: str,
    task: str,
) -> Any:
    """Instantiate the appropriate ActionTracker variant.

    Falls back to :class:`_NoOpTracker` when the ``matrice`` package is
    not installed, so the Trainer can run in local / CI environments.

    Args:
        action_id: Platform action ID, or ``None`` for local testing.
        architecture: Model architecture name (for the local tracker).
        task: Task type string (for the local tracker).

    Returns:
        An ``ActionTracker``, ``LocalActionTracker``, or
        :class:`_NoOpTracker`.
    """
    logger.debug(
        "Creating action tracker for architecture=%s task=%s",
        architecture,
        task,
    )
    try:
        if action_id is not None:
            from matrice.actionTracker import ActionTracker

            return ActionTracker(action_id)

        # action_id is None: use no-op for local/CI runs (avoids LocalActionTracker
        # prompts or platform dependencies)
        return _NoOpTracker()
    except ImportError:
        logger.warning(
            "matrice.actionTracker not available — using no-op tracker. "
            "Install the 'matrice' package for full platform integration.",
        )
        return _NoOpTracker()


# ------------------------------------------------------------------ #
#  No-op ActionTracker stub                                            #
# ------------------------------------------------------------------ #


class _NoOpTracker:
    """Stub ActionTracker for environments without the ``matrice`` SDK.

    Every method is a silent no-op so that the Trainer lifecycle can
    run without side-effects during local development or CI.
    """

    def update_status(
        self,
        _code: str,
        _level: str,
        _message: str,
        /,
    ) -> None:
        """Accept and discard a status update."""

    def log_error(
        self,
        _file: str,
        _location: str,
        _message: str,
        /,
    ) -> None:
        """Accept and discard an error log entry."""

    def log_epoch_results(
        self,
        _epoch: int,
        _metrics: list[dict[str, Any]],
        /,
    ) -> None:
        """Accept and discard epoch metrics."""

    def add_index_to_category(
        self,
        _mapping: dict[str, str],
        /,
    ) -> None:
        """Accept and discard a label-index mapping."""

    def upload_checkpoint(self, _path: str, /) -> None:
        """Accept and discard a checkpoint upload request."""

    def save_evaluation_results(
        self,
        _payload: list[dict[str, Any]],
        /,
    ) -> None:
        """Accept and discard evaluation results."""

    def get_job_params(self) -> dict[str, Any]:
        """Return empty job parameters."""
        return {}

    def get_checkpoint_path(
        self,
        _config: Any = None,
        /,
    ) -> tuple[str | None, bool]:
        """Return no-checkpoint sentinel."""
        return (None, False)
