"""Atomic checkpoints preserving optimizer, sampler, RNG and selected adapter."""

import json
import random
import shutil
import uuid
from pathlib import Path

import numpy as np
import torch

from .common import atomic_json, sha_file


def rng_state():
    return {
        "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,
    }


def restore_rng(state):
    random.setstate(state["python"])
    np.random.set_state(state["numpy"])
    torch.set_rng_state(state["torch"].cpu())
    if state["cuda"] is not None and torch.cuda.is_available():
        torch.cuda.set_rng_state_all([s.cpu() for s in state["cuda"]])


def save_checkpoint(run_dir, model, optimizer, scheduler, state, cfg, codec, commit=None):
    run_dir = Path(run_dir)
    folder = run_dir / "checkpoints"
    folder.mkdir(parents=True, exist_ok=True)
    temporary = folder / (".writing-" + uuid.uuid4().hex)
    temporary.mkdir()
    model.save_adapter(temporary / "adapter", codec, cfg)
    torch.save(
        {
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict(),
            "rng": rng_state(),
            "state": state,
        },
        temporary / "training.pt",
    )
    atomic_json(
        temporary / "complete.json",
        {
            "step": state["step"],
            "training_sha256": sha_file(temporary / "training.pt"),
            "files": {
                str(p.relative_to(temporary)): sha_file(p)
                for p in temporary.rglob("*")
                if p.is_file()
            },
        },
    )
    target = folder / f"step-{state['step']:07d}-{uuid.uuid4().hex[:6]}"
    temporary.rename(target)
    previous = (
        json.loads((run_dir / "latest.json").read_text())
        if (run_dir / "latest.json").exists()
        else None
    )
    atomic_json(run_dir / "latest.json", {"path": str(target.relative_to(run_dir))})
    if previous:
        atomic_json(run_dir / "previous.json", previous)
    if commit:
        commit()
    retain = set()
    for marker in ("latest.json", "previous.json", "best.json"):
        if (run_dir / marker).exists():
            retain.add(json.loads((run_dir / marker).read_text())["path"])
    for path in folder.glob("step-*"):
        if str(path.relative_to(run_dir)) not in retain:
            shutil.rmtree(path)
    return target


def select_best(run_dir, checkpoint, loss, commit=None):
    run_dir = Path(run_dir)
    atomic_json(
        run_dir / "best.json", {"path": str(Path(checkpoint).relative_to(run_dir)), "loss": loss}
    )
    if commit:
        commit()


def checkpoint_path(run_dir, marker="latest"):
    root = Path(run_dir)
    pointer = json.loads((root / (marker + ".json")).read_text())
    path = root / pointer["path"]
    if not path.resolve().is_relative_to(root.resolve()):
        raise ValueError("Checkpoint pointer escapes run directory")
    manifest = json.loads((path / "complete.json").read_text())
    if sha_file(path / "training.pt") != manifest["training_sha256"]:
        raise ValueError("Incomplete or damaged checkpoint")
    for name, expected in manifest.get("files", {}).items():
        candidate = path / name
        if (
            not candidate.resolve().is_relative_to(path.resolve())
            or sha_file(candidate) != expected
        ):
            raise ValueError("Incomplete or damaged checkpoint adapter")
    return path


def load_training(path, optimizer, scheduler):
    # These files are generated locally by this run, never loaded from arbitrary Hub repos.
    saved = torch.load(Path(path) / "training.pt", map_location="cpu", weights_only=False)
    optimizer.load_state_dict(saved["optimizer"])
    scheduler.load_state_dict(saved["scheduler"])
    restore_rng(saved["rng"])
    return saved["state"]