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

import torch
from safetensors.torch import load_model, save_model


def file_hash(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def save_checkpoint(path: Path, model, optimizer, scheduler, progress: dict, identity: str):
    if path.exists():
        raise FileExistsError(f"Refusing to overwrite checkpoint: {path}")
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.parent / f".checkpoint-{uuid.uuid4().hex}"
    temporary.mkdir()
    try:
        save_model(model, temporary / "model.safetensors")
        state = {
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict(),
            "python_rng": random.getstate(),
            "torch_rng": torch.get_rng_state(),
            "cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [],
        }
        torch.save(state, temporary / "training.pt")
        (temporary / "progress.json").write_text(json.dumps(progress, indent=2) + "\n")
        marker = {
            "identity": identity,
            "files": {
                name: file_hash(temporary / name)
                for name in ("model.safetensors", "training.pt", "progress.json")
            },
        }
        (temporary / "complete.json").write_text(json.dumps(marker, indent=2) + "\n")
        os.replace(temporary, path)
    except BaseException:
        shutil.rmtree(temporary)
        raise


def verify_checkpoint(path: Path, identity: str):
    marker = json.loads((path / "complete.json").read_text())
    if marker["identity"] != identity:
        raise ValueError("Checkpoint protocol identity mismatch")
    expected = {"model.safetensors", "training.pt", "progress.json"}
    if set(marker["files"]) != expected:
        raise ValueError("Checkpoint file inventory mismatch")
    for name, digest in marker["files"].items():
        if file_hash(path / name) != digest:
            raise ValueError(f"Checkpoint checksum mismatch: {name}")


def load_checkpoint(path: Path, model, optimizer, scheduler, identity: str) -> dict:
    verify_checkpoint(path, identity)
    load_model(model, path / "model.safetensors", strict=True)
    state = torch.load(path / "training.pt", map_location="cpu", weights_only=True)
    optimizer.load_state_dict(state["optimizer"])
    scheduler.load_state_dict(state["scheduler"])
    random.setstate(state["python_rng"])
    torch.set_rng_state(state["torch_rng"])
    if state["cuda_rng"]:
        if not torch.cuda.is_available() or len(state["cuda_rng"]) != torch.cuda.device_count():
            raise ValueError("CUDA random state cannot be restored on this worker")
        torch.cuda.set_rng_state_all(state["cuda_rng"])
    return json.loads((path / "progress.json").read_text())