"""Atomic training checkpoints with full RNG state for exact resume."""

from __future__ import annotations

import os
import random
import re
from pathlib import Path

import numpy as np
import torch

_STEP_RE = re.compile(r"step_(\d{7})\.pt$")


def checkpoint_path(directory: Path, step: int) -> Path:
    return Path(directory) / f"step_{step:07d}.pt"


def save_checkpoint(path: Path, model, optimizer, scaler, step: int, extra: dict) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    state = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict() if optimizer is not None else None,
        "scaler": scaler.state_dict() if scaler is not None else None,
        "step": step,
        "extra": extra,
        "rng": {
            "python": random.getstate(),
            "numpy": np.random.get_state(),
            "torch": torch.get_rng_state(),
            "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
        },
    }
    tmp = path.with_suffix(".tmp")
    torch.save(state, tmp)
    os.replace(tmp, path)


def load_checkpoint(path: Path, model, optimizer=None, scaler=None) -> dict:
    state = torch.load(path, map_location="cpu", weights_only=False)
    model.load_state_dict(state["model"])
    if optimizer is not None and state["optimizer"] is not None:
        optimizer.load_state_dict(state["optimizer"])
    if scaler is not None and state["scaler"] is not None:
        scaler.load_state_dict(state["scaler"])
    rng = state["rng"]
    random.setstate(rng["python"])
    np.random.set_state(rng["numpy"])
    torch.set_rng_state(rng["torch"])
    if rng["cuda"] is not None and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(rng["cuda"])
    return {"step": state["step"], "extra": state["extra"]}


def _step_files(directory: Path) -> list[tuple[int, Path]]:
    out = []
    for p in Path(directory).glob("step_*.pt"):
        m = _STEP_RE.search(p.name)
        if m:
            out.append((int(m.group(1)), p))
    return sorted(out)


def latest_checkpoint(directory: Path) -> Path | None:
    files = _step_files(directory) if Path(directory).exists() else []
    return files[-1][1] if files else None


def prune_checkpoints(directory: Path, keep_last: int = 2) -> None:
    for _, p in _step_files(directory)[:-keep_last]:
        p.unlink()