"""Callback system for the Trainer lifecycle.

Defines :class:`TrainerCallback` (base with no-op defaults) and three
built-in callbacks that cover the most common cross-cutting concerns:

* :class:`ActionTrackerCallback` — epoch logging and status updates.
* :class:`EarlyStoppingCallback` — patience-based early stopping.
* :class:`CheckpointCallback` — periodic model/state checkpointing.

Custom callbacks subclass :class:`TrainerCallback` and override only
the hooks they need.
"""

from __future__ import annotations

import logging
import math
import shutil
from pathlib import Path
from typing import TYPE_CHECKING, Any


if TYPE_CHECKING:
    from matrice_models.training.trainer import Trainer, TrainResult


logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
#  Base callback                                                       #
# ------------------------------------------------------------------ #


class TrainerCallback:
    """Base callback with no-op defaults for every hook.

    Subclass and override only the methods you need.  Any method not
    overridden silently does nothing, so new hooks can be added in
    future releases without breaking existing callbacks.

    Hook firing order inside ``LoopTrainer.train_loop``::

        on_train_begin(trainer)
        for epoch in range(epochs):
            on_epoch_begin(trainer, epoch)
            ... train + validate ...
            on_epoch_end(trainer, epoch, metrics)
        on_train_end(trainer, result)
    """

    def on_train_begin(self, trainer: Trainer) -> None:
        """Called once before the first epoch.

        Args:
            trainer: The active ``Trainer`` instance.
        """

    def on_epoch_begin(self, trainer: Trainer, epoch: int) -> None:
        """Called at the start of each epoch.

        Args:
            trainer: The active ``Trainer`` instance.
            epoch: Zero-based epoch index.
        """

    def on_epoch_end(
        self,
        trainer: Trainer,
        epoch: int,
        metrics: list[dict[str, Any]],
    ) -> None:
        """Called after training *and* validation for one epoch complete.

        Args:
            trainer: The active ``Trainer`` instance.
            epoch: Zero-based epoch index.
            metrics: Epoch metrics in ActionTracker format — a list of
                dicts each containing ``splitType``, ``metricName``, and
                ``metricValue``.
        """

    def on_train_end(self, trainer: Trainer, result: TrainResult) -> None:
        """Called once after the last epoch (or early-stop).

        Args:
            trainer: The active ``Trainer`` instance.
            result: The :class:`TrainResult` produced by the training loop.
        """


# ------------------------------------------------------------------ #
#  Built-in: ActionTracker epoch logger                                #
# ------------------------------------------------------------------ #


class ActionTrackerCallback(TrainerCallback):
    """Forward per-epoch metrics and key status updates to the ActionTracker.

    Registered by default in ``LoopTrainer`` so that every epoch's
    train/val metrics are recorded on the platform without any extra
    code in the BYOM project.
    """

    def on_train_begin(self, trainer: Trainer) -> None:
        """Send the ``MDL_TRN_STR`` status to signal loop start."""
        trainer.action_tracker.update_status(
            "MDL_TRN_STR",
            "OK",
            "Model training is starting",
        )

    def on_epoch_end(
        self,
        trainer: Trainer,
        epoch: int,
        metrics: list[dict[str, Any]],
    ) -> None:
        """Log epoch metrics via ``action_tracker.log_epoch_results``."""
        from matrice_models.training.utils import sanitise_metrics

        sanitised = sanitise_metrics(metrics)
        trainer.action_tracker.log_epoch_results(epoch, sanitised)


# ------------------------------------------------------------------ #
#  Built-in: Early stopping                                            #
# ------------------------------------------------------------------ #


class EarlyStoppingCallback(TrainerCallback):
    """Stop training when a monitored metric stops improving.

    Tracks the value of *monitor* across epochs.  If the metric does
    not improve by at least *min_delta* for *patience* consecutive
    epochs, :attr:`should_stop` is set to ``True`` and the training
    loop should break.

    Args:
        patience: Number of epochs to wait without improvement.
        min_delta: Minimum change to qualify as an improvement.
        monitor: Metric name to watch (matched against the
            ``metricName`` key in the epoch-metrics list).
        mode: ``"min"`` treats lower values as better (e.g. loss),
            ``"max"`` treats higher values as better (e.g. accuracy).
    """

    def __init__(
        self,
        patience: int = 5,
        min_delta: float = 0.0,
        monitor: str = "loss",
        mode: str = "min",
    ) -> None:
        """Initialize early-stopping state.

        Args:
            patience: Epochs without improvement before stopping.
            min_delta: Minimum absolute improvement threshold.
            monitor: ``metricName`` string to track in epoch metrics.
            mode: ``"min"`` (loss-like) or ``"max"`` (accuracy-like).
        """
        if mode not in ("min", "max"):
            msg = f"mode must be 'min' or 'max', got '{mode}'"
            raise ValueError(msg)

        self.patience = patience
        self.min_delta = abs(min_delta)
        self.monitor = monitor
        self.mode = mode

        self.best_value: float = math.inf if mode == "min" else -math.inf
        self.wait_count: int = 0
        self.best_epoch: int = 0
        self.should_stop: bool = False

    def on_epoch_end(
        self,
        trainer: Trainer,
        epoch: int,
        metrics: list[dict[str, Any]],
    ) -> None:
        """Check the monitored metric and update the patience counter."""
        current = self._extract_metric(metrics)
        if current is None:
            logger.warning(
                "EarlyStoppingCallback: metric '%s' not found in epoch %d — skipping check",
                self.monitor,
                epoch,
            )
            return

        if self._is_improvement(current):
            self.best_value = current
            self.wait_count = 0
            self.best_epoch = epoch
        else:
            self.wait_count += 1
            logger.info(
                "EarlyStopping: no improvement for %d/%d epochs (best=%.6f at epoch %d)",
                self.wait_count,
                self.patience,
                self.best_value,
                self.best_epoch,
            )
            if self.wait_count >= self.patience:
                self.should_stop = True
                logger.info("EarlyStopping triggered after %d epochs without improvement", self.patience)

    def _is_improvement(self, current: float) -> bool:
        """Return ``True`` if *current* is better than *best_value* by at least *min_delta*."""
        if self.mode == "min":
            return current < self.best_value - self.min_delta
        return current > self.best_value + self.min_delta

    def _extract_metric(
        self,
        metrics: list[dict[str, Any]],
    ) -> float | None:
        """Pull the monitored value from the metrics list.

        Prefers ``splitType == "val"``; falls back to any split.
        """
        val_match: float | None = None
        any_match: float | None = None

        for m in metrics:
            if m.get("metricName") == self.monitor:
                value = m.get("metricValue")
                if value is None:
                    continue
                any_match = float(value)
                if m.get("splitType") == "val":
                    val_match = float(value)

        return val_match if val_match is not None else any_match


# ------------------------------------------------------------------ #
#  Built-in: Checkpoint saver                                          #
# ------------------------------------------------------------------ #


class CheckpointCallback(TrainerCallback):
    """Save model checkpoints periodically during training.

    By default a checkpoint is written after every epoch that improves
    the monitored metric.  Set ``save_every_n_epochs`` to checkpoint at
    a fixed cadence regardless of metric changes.

    Args:
        checkpoint_dir: Directory where checkpoint files are written.
        monitor: Metric name that determines "best" (matched against
            ``metricName`` in epoch-metrics).
        mode: ``"min"`` or ``"max"`` — how *monitor* is compared.
        save_every_n_epochs: If set, save a checkpoint every *N* epochs
            regardless of metric improvement.  ``None`` means save only
            when the monitored metric improves.
    """

    def __init__(
        self,
        checkpoint_dir: str | Path = "checkpoints",
        monitor: str = "loss",
        mode: str = "min",
        save_every_n_epochs: int | None = None,
    ) -> None:
        """Initialize checkpoint state.

        Args:
            checkpoint_dir: Target directory for checkpoint files.
            monitor: ``metricName`` string to track for best-model logic.
            mode: ``"min"`` (loss-like) or ``"max"`` (accuracy-like).
            save_every_n_epochs: Fixed checkpoint cadence, or ``None``
                for best-only saving.
        """
        if mode not in ("min", "max"):
            msg = f"mode must be 'min' or 'max', got '{mode}'"
            raise ValueError(msg)

        self.checkpoint_dir = Path(checkpoint_dir)
        self.monitor = monitor
        self.mode = mode
        self.save_every_n_epochs = save_every_n_epochs

        self.best_value: float = math.inf if mode == "min" else -math.inf

    def on_epoch_end(
        self,
        trainer: Trainer,
        epoch: int,
        metrics: list[dict[str, Any]],
    ) -> None:
        """Save a checkpoint if conditions are met."""
        periodic_save = self.save_every_n_epochs is not None and (epoch + 1) % self.save_every_n_epochs == 0

        current = self._extract_metric(metrics)
        is_best = current is not None and self._is_improvement(current)

        if is_best:
            self.best_value = current  # type: ignore[assignment]

        if is_best or periodic_save:
            self._save(trainer, epoch, is_best=is_best)

    def _is_improvement(self, current: float) -> bool:
        """Return ``True`` if *current* beats the stored best."""
        if self.mode == "min":
            return current < self.best_value
        return current > self.best_value

    def _save(
        self,
        trainer: Trainer,
        epoch: int,
        *,
        is_best: bool,
    ) -> None:
        """Persist the checkpoint to disk and upload via ActionTracker."""
        import torch

        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        filename = f"checkpoint_epoch_{epoch}.pth.tar"
        filepath = self.checkpoint_dir / filename

        state: dict[str, Any] = {
            "epoch": epoch,
            "architecture": trainer.model_info.architecture,
        }

        model = trainer.model
        if model is not None and hasattr(model, "state_dict"):
            state["state_dict"] = model.state_dict()

        torch.save(state, str(filepath))
        logger.info("Checkpoint saved to %s", filepath)

        if is_best:
            best_path = self.checkpoint_dir / "model_best.pth.tar"
            shutil.copy2(filepath, best_path)
            logger.info("Best checkpoint copied to %s", best_path)
            trainer.action_tracker.upload_checkpoint(str(best_path))

    def _extract_metric(
        self,
        metrics: list[dict[str, Any]],
    ) -> float | None:
        """Pull the monitored value from the metrics list (prefers val split)."""
        val_match: float | None = None
        any_match: float | None = None

        for m in metrics:
            if m.get("metricName") == self.monitor:
                value = m.get("metricValue")
                if value is None:
                    continue
                any_match = float(value)
                if m.get("splitType") == "val":
                    val_match = float(value)

        return val_match if val_match is not None else any_match


# ------------------------------------------------------------------ #
#  Helpers                                                             #
# ------------------------------------------------------------------ #


def run_callbacks(
    callbacks: list[TrainerCallback],
    hook: str,
    *args: Any,
    **kwargs: Any,
) -> None:
    """Fire a named hook on every callback in the list.

    Catches and logs exceptions from individual callbacks so that a
    broken callback does not crash the training loop.

    Args:
        callbacks: List of :class:`TrainerCallback` instances.
        hook: Method name to invoke (e.g. ``"on_epoch_end"``).
        *args: Positional arguments forwarded to the hook.
        **kwargs: Keyword arguments forwarded to the hook.
    """
    for cb in callbacks:
        method = getattr(cb, hook, None)
        if method is None:
            continue
        try:
            method(*args, **kwargs)
        except Exception:
            logger.exception(
                "Callback %s.%s raised an exception — continuing",
                type(cb).__name__,
                hook,
            )
