"""Checkpoint caching module for scDistill.

Provides configuration and utilities for saving model checkpoints
during training at specified epoch intervals.

Example
-------
>>> from scdistill import Distiller, CacheConfig
>>>
>>> cache_config = CacheConfig(
...     cache_dir="./checkpoints",
...     save_every_n_epochs=10,
...     save_latent=True,
... )
>>> model = Distiller(adata, batch_key='batch')
>>> model.train(cache_config=cache_config)
"""

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Union, Literal

import torch
import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class CacheConfig:
    """Configuration for checkpoint caching during training.

    Parameters
    ----------
    cache_dir : str or Path
        Base directory for storing checkpoints.
        Creates subdirectories: cache_dir/epoch_{N}/
    save_every_n_epochs : int, default=10
        Interval (in epochs) for saving checkpoints.
        Set to 1 to save every epoch (caution: high disk usage).
    save_encoder : bool, default=True
        If True, save encoder.pt at each checkpoint.
    save_decoder : bool, default=True
        If True, save decoder.pt at each checkpoint.
    save_latent : bool, default=True
        If True, save latent.npy (Z from encoder) at each checkpoint.
        Note: During encoder training, this is the encoder's output.
        During decoder training, latent is not re-saved (encoder is frozen).

    Examples
    --------
    >>> cache_config = CacheConfig(
    ...     cache_dir="./checkpoints",
    ...     save_every_n_epochs=10,
    ... )
    >>> model.train(cache_config=cache_config)
    """

    cache_dir: Union[str, Path]
    save_every_n_epochs: int = 10
    save_encoder: bool = True
    save_decoder: bool = True
    save_latent: bool = True
    # Internal: set by training functions
    _phase: Optional[Literal["encoder", "decoder"]] = field(default=None, repr=False)

    def __post_init__(self):
        """Validate configuration."""
        self.cache_dir = Path(self.cache_dir)
        if self.save_every_n_epochs < 1:
            raise ValueError("save_every_n_epochs must be >= 1")

    def get_checkpoint_dir(self, epoch: int) -> Path:
        """Get checkpoint directory for a specific epoch.

        Parameters
        ----------
        epoch : int
            Epoch number (0-indexed).

        Returns
        -------
        Path
            Directory path for checkpoint.
        """
        if self._phase:
            return self.cache_dir / self._phase / f"epoch_{epoch + 1:04d}"
        return self.cache_dir / f"epoch_{epoch + 1:04d}"

    def should_save(self, epoch: int) -> bool:
        """Determine if checkpoint should be saved at this epoch.

        Parameters
        ----------
        epoch : int
            Current epoch (0-indexed).

        Returns
        -------
        bool
            True if checkpoint should be saved.
        """
        # epoch is 0-indexed, so epoch 9 is the 10th epoch
        # Save at epochs: 9, 19, 29, ... (for save_every_n_epochs=10)
        return (epoch + 1) % self.save_every_n_epochs == 0


def save_checkpoint(
    cache_config: CacheConfig,
    epoch: int,
    encoder: Optional[torch.nn.Module] = None,
    decoder: Optional[torch.nn.Module] = None,
    latent: Optional[np.ndarray] = None,
    loss: Optional[float] = None,
) -> Path:
    """Save a training checkpoint.

    Parameters
    ----------
    cache_config : CacheConfig
        Checkpoint configuration.
    epoch : int
        Current epoch number (0-indexed).
    encoder : nn.Module, optional
        Encoder model to save.
    decoder : nn.Module, optional
        Decoder model to save.
    latent : np.ndarray, optional
        Latent representation to save.
    loss : float, optional
        Current loss value (saved to metadata).

    Returns
    -------
    Path
        Path to checkpoint directory.
    """
    checkpoint_dir = cache_config.get_checkpoint_dir(epoch)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    # Save encoder
    if cache_config.save_encoder and encoder is not None:
        encoder_path = checkpoint_dir / "encoder.pt"
        torch.save(encoder.state_dict(), encoder_path)

    # Save decoder
    if cache_config.save_decoder and decoder is not None:
        decoder_path = checkpoint_dir / "decoder.pt"
        torch.save(decoder.state_dict(), decoder_path)

    # Save latent representation
    if cache_config.save_latent and latent is not None:
        latent_path = checkpoint_dir / "latent.npy"
        np.save(latent_path, latent)

    # Save metadata
    metadata = {
        "epoch": epoch + 1,  # 1-indexed for display
        "phase": cache_config._phase,
        "loss": loss,
    }
    metadata_path = checkpoint_dir / "metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)

    logger.info(f"  Checkpoint saved: {checkpoint_dir}")

    return checkpoint_dir
