"""Delegated trainer for framework-owned training loops.

``DelegatedTrainer`` extends :class:`~matrice_models.training.trainer.Trainer`
for models whose training loop is owned by an **external framework** — the
trainer simply configures and launches the framework, then collects results.

This is the pattern used by:

* **Ultralytics YOLO** — ``model.train(data=..., epochs=...)``
* **D-FINE solver** — ``solver.fit()``

Subclasses implement :meth:`launch_training`, which calls the framework's
train API and returns a :class:`~matrice_models.training.trainer.TrainResult`.
All ActionTracker lifecycle phases, error handling, and checkpoint upload are
still managed by the base :class:`Trainer`.

Example (Ultralytics YOLO)::

    class YOLOTrainer(DelegatedTrainer):
        def setup_data(self) -> None:
            self.data_yaml = f"{self.config.dataset_path}/data.yaml"

        def setup_model(self) -> None:
            from ultralytics import YOLO

            self.model = YOLO(f"{self.model_info.architecture}.pt")

        def get_label_map(self) -> dict[str, str]:
            return self.model.names

        def launch_training(self) -> TrainResult:
            results = self.model.train(
                data=self.data_yaml,
                epochs=self.config.epochs,
                batch=self.config.batch_size,
            )
            return TrainResult(
                best_metric_name="mAP50",
                best_metric_value=results.results_dict["metrics/mAP50(B)"],
                best_epoch=0,
                epochs_completed=self.config.epochs,
            )

        def get_best_checkpoint_path(self) -> Path | None:
            return Path("runs/detect/train/weights/best.pt")

        def run_evaluation(self) -> list[dict[str, Any]]: ...


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

from __future__ import annotations

import logging
from abc import abstractmethod
from typing import TYPE_CHECKING

from matrice_models.training.trainer import Trainer, TrainResult


if TYPE_CHECKING:
    from pathlib import Path

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


logger = logging.getLogger(__name__)


class DelegatedTrainer(Trainer):
    """Trainer for framework-owned training loops.

    Unlike :class:`~matrice_models.training.loop_trainer.LoopTrainer`,
    which manages the epoch loop itself, ``DelegatedTrainer`` hands
    control entirely to the external framework (Ultralytics, D-FINE
    solver, etc.).  Subclasses implement:

    ========================  =============================================
    Hook                      Responsibility
    ========================  =============================================
    launch_training           Call the framework's train API, return result
    get_best_checkpoint_path  Return the path to the best model file
    ========================  =============================================

    The remaining hooks (``setup_data``, ``setup_model``,
    ``get_label_map``, ``run_evaluation``) are inherited from
    :class:`Trainer`.

    Callback support is intentionally limited — frameworks like
    Ultralytics have their **own** callback system.  The
    ``on_train_begin`` / ``on_train_end`` hooks still fire, but
    per-epoch hooks (``on_epoch_begin``, ``on_epoch_end``) should
    be wired through the framework's native callbacks instead.
    """

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

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

        Args:
            action_id: Platform action ID (``None`` for local runs).
            model_info: Model identity and capabilities.
            config: Training hyperparameters.
            checkpoint_dir: Directory for checkpoint files (note:
                the framework may write checkpoints to its own
                location — see :meth:`get_best_checkpoint_path`).
            callbacks: Optional list of callbacks.
        """
        super().__init__(
            action_id,
            model_info,
            config,
            checkpoint_dir=checkpoint_dir,
            callbacks=callbacks,
        )

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

    @abstractmethod
    def launch_training(self) -> TrainResult:
        """Call the framework's training API and return results.

        This is where the framework takes full control.  The
        implementation should:

        1. Call the framework's train method (e.g. ``model.train()``,
           ``solver.fit()``).
        2. Collect final metrics from the framework's results.
        3. Build and return a :class:`TrainResult`.

        The base trainer handles ActionTracker status updates *around*
        this call — no need to send ``MDL_TRN_STR`` / ``MDL_TRN_CMPL``
        manually.

        Returns:
            A :class:`TrainResult` populated with the training outcome.

        Example (Ultralytics)::

            def launch_training(self) -> TrainResult:
                results = self.model.train(
                    data=self.data_yaml,
                    epochs=self.config.epochs,
                )
                return TrainResult(
                    best_metric_name="mAP50",
                    best_metric_value=results.results_dict["metrics/mAP50(B)"],
                    best_epoch=0,
                    epochs_completed=self.config.epochs,
                )
        """

    @abstractmethod
    def get_best_checkpoint_path(self) -> Path | None:
        """Return the path to the framework's best model checkpoint.

        Frameworks typically save their own checkpoints in a
        framework-specific location (e.g. Ultralytics writes to
        ``runs/detect/train/weights/best.pt``).  This method tells
        the trainer where to find that file for upload.

        Returns:
            Path to the best checkpoint file, or ``None`` if not
            available.

        Example::

            def get_best_checkpoint_path(self) -> Path:
                return Path("runs/detect/train/weights/best.pt")
        """

    # ------------------------------------------------------------------ #
    #  Overridable hooks                                                   #
    # ------------------------------------------------------------------ #

    def register_framework_callbacks(self) -> None:
        """Register callbacks with the framework's own callback system.

        Override this to wire framework-native callbacks for per-epoch
        logging.  Called once before :meth:`launch_training`.

        Example (Ultralytics)::

            def register_framework_callbacks(self) -> None:
                self.model.add_callback("on_fit_epoch_end", self._log_yolo_epoch)
        """

    # ------------------------------------------------------------------ #
    #  train_loop implementation                                           #
    # ------------------------------------------------------------------ #

    def train_loop(self) -> None:
        """Delegate training to the framework and collect results.

        Lifecycle::

            register_framework_callbacks()
            on_train_begin
            launch_training()  →  TrainResult
            on_train_end(result)

        Populates ``self._result`` before returning.
        """
        self.register_framework_callbacks()
        self.notify("on_train_begin", self)

        result = self.launch_training()
        self._result = result

        self.notify("on_train_end", self, result)

    def save_model(self) -> None:
        """Upload the framework's best checkpoint to the platform.

        Uses :meth:`get_best_checkpoint_path` to locate the file.
        Falls back to the base class implementation if no path is
        returned.
        """
        best_path = self.get_best_checkpoint_path()
        if best_path is not None and best_path.exists():
            self.action_tracker.upload_checkpoint(str(best_path))
            logger.info("Uploaded framework checkpoint: %s", best_path)
        else:
            super().save_model()
