"""Checkpoint save / load utilities for PyTorch training.

Extracted from the ``save_checkpoint()`` functions duplicated in
``ml_pytorch_vision_classification/train.py`` and ``r2plus1d/train.py``.
"""

from __future__ import annotations

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


if TYPE_CHECKING:
    import torch.nn as nn


logger = logging.getLogger(__name__)


def save_checkpoint(
    state: dict[str, Any],
    model: nn.Module | Any,
    *,
    is_best: bool,
    checkpoint_dir: str | Path = "checkpoints",
    filename: str = "checkpoint.pth.tar",
) -> list[Path]:
    """Save a training checkpoint to disk.

    Writes ``state`` (typically epoch, state_dict, optimizer, scheduler)
    to *checkpoint_dir/filename*.  When *is_best* is ``True``, also:

    * Copies the checkpoint to ``model_best.pth.tar``.
    * Saves the raw model to ``model_best.pt`` (handles
      ``DataParallel`` unwrapping automatically).

    Args:
        state: Serialisable dict — must include at minimum ``epoch``
            and ``state_dict``.
        model: The PyTorch model (or ``DataParallel``-wrapped model).
            Used only when *is_best* to create the ``.pt`` snapshot.
        is_best: Whether the current checkpoint is the best so far.
        checkpoint_dir: Directory where files are written.
        filename: Name of the periodic checkpoint file.

    Returns:
        List of :class:`Path` objects for every file written.
    """
    import torch

    checkpoint_dir = Path(checkpoint_dir)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    filepath = checkpoint_dir / filename
    torch.save(state, str(filepath))
    logger.info("Checkpoint saved to %s", filepath)
    written: list[Path] = [filepath]

    if is_best:
        best_pth = checkpoint_dir / "model_best.pth.tar"
        shutil.copyfile(filepath, best_pth)
        logger.info("Best checkpoint saved to %s", best_pth)
        written.append(best_pth)

        raw_model = _unwrap_model(model)
        best_pt = checkpoint_dir / "model_best.pt"
        torch.save(raw_model, str(best_pt))
        logger.info("Best model (PT format) saved to %s", best_pt)
        written.append(best_pt)

    return written


def load_checkpoint(
    path: str | Path,
    *,
    map_location: str = "cpu",
) -> dict[str, Any]:
    """Load a checkpoint from disk.

    Args:
        path: Path to the ``.pth.tar`` checkpoint file.
        map_location: Device to map tensors to (default ``"cpu"``
            for safe cross-device loading).

    Returns:
        The checkpoint state dict.

    Raises:
        FileNotFoundError: If *path* does not exist.
    """
    import torch

    path = Path(path)
    if not path.exists():
        msg = f"Checkpoint not found: {path}"
        raise FileNotFoundError(msg)

    logger.info("Loading checkpoint from %s", path)
    return torch.load(str(path), map_location=map_location, weights_only=False)


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


def _unwrap_model(model: nn.Module | Any) -> nn.Module | Any:
    """Unwrap ``DataParallel`` / ``DistributedDataParallel`` if needed.

    Also handles wrapper objects that store the real model in a
    ``.model`` attribute (e.g. r2plus1d_Model).
    """
    import torch.nn as nn_module

    if hasattr(model, "model"):
        model = model.model

    if isinstance(model, (nn_module.DataParallel, nn_module.parallel.DistributedDataParallel)):
        return copy.deepcopy(model.module)

    return model
