"""Model checkpoint management for SalTTS."""

import shutil
from pathlib import Path
from typing import Any, Dict, Optional

import torch


class CheckpointManager:
    """Manages model checkpoints with automatic cleanup."""

    def __init__(self, checkpoint_dir: str, keep_last_n: int = 3, save_best: bool = True):
        """Initialize checkpoint manager.

        Args:
            checkpoint_dir: Directory to save checkpoints
            keep_last_n: Number of recent checkpoints to keep
            save_best: Whether to save best checkpoint separately
        """
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.keep_last_n = keep_last_n
        self.save_best = save_best
        self.best_metric: Optional[float] = None

    def save(
        self,
        model: torch.nn.Module,
        optimizer: Optional[torch.optim.Optimizer],
        epoch: int,
        metrics: Optional[Dict[str, float]] = None,
        extra_data: Optional[Dict[str, Any]] = None,
    ) -> Path:
        """Save a checkpoint.

        Args:
            model: Model to save
            optimizer: Optimizer to save
            epoch: Current epoch
            metrics: Training metrics
            extra_data: Additional data to save

        Returns:
            Path to saved checkpoint
        """
        checkpoint = {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "metrics": metrics or {},
        }

        if optimizer:
            checkpoint["optimizer_state_dict"] = optimizer.state_dict()

        if extra_data:
            checkpoint.update(extra_data)

        # Save regular checkpoint
        checkpoint_path = self.checkpoint_dir / f"checkpoint_epoch_{epoch}.pt"
        torch.save(checkpoint, checkpoint_path)

        # Save best checkpoint if applicable
        if self.save_best and metrics:
            metric_value = metrics.get("val_loss", float("inf"))
            if self.best_metric is None or metric_value < self.best_metric:
                self.best_metric = metric_value
                best_path = self.checkpoint_dir / "best_checkpoint.pt"
                shutil.copy(checkpoint_path, best_path)

        # Cleanup old checkpoints
        self._cleanup_old_checkpoints()

        return checkpoint_path

    def load(self, checkpoint_path: str, model: torch.nn.Module, optimizer: Optional[torch.optim.Optimizer] = None) -> Dict[str, Any]:
        """Load a checkpoint.

        Args:
            checkpoint_path: Path to checkpoint file
            model: Model to load state into
            optimizer: Optimizer to load state into

        Returns:
            Checkpoint data
        """
        checkpoint = torch.load(checkpoint_path, map_location="cpu")
        model.load_state_dict(checkpoint["model_state_dict"])

        if optimizer and "optimizer_state_dict" in checkpoint:
            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

        return checkpoint

    def _cleanup_old_checkpoints(self) -> None:
        """Remove old checkpoints, keeping only the last N."""
        checkpoints = sorted(
            self.checkpoint_dir.glob("checkpoint_epoch_*.pt"),
            key=lambda p: p.stat().st_mtime,
        )

        if len(checkpoints) > self.keep_last_n:
            for checkpoint in checkpoints[: -self.keep_last_n]:
                checkpoint.unlink()
