"""Checkpointing: resumable experiments, never rely on RAM."""
from __future__ import annotations

import json
import os
import time
from dataclasses import asdict

import torch

from .model import TinyChess, TinyChessConfig


def save_checkpoint(path, model, cfg, optimizer=None, scheduler=None, step=0,
                    metrics=None, extra=None, rng=True):
    tmp = path + ".tmp"
    payload = {
        "model": model.state_dict(),
        "cfg": cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg),
        "step": step,
        "metrics": metrics or {},
        "extra": extra or {},
        "param_report": model.param_report(),
        "saved_at": time.time(),
    }
    if optimizer is not None:
        payload["optimizer"] = optimizer.state_dict()
    if scheduler is not None:
        payload["scheduler"] = scheduler.state_dict()
    if rng:
        payload["rng"] = {
            "torch": torch.get_rng_state(),
            "python": __import__("random").getstate(),
            "numpy": __import__("numpy").random.get_state(),
        }
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    torch.save(payload, tmp)
    os.replace(tmp, path)
    return path


def load_checkpoint(path, map_location="cpu", restore_rng=True):
    ck = torch.load(path, map_location=map_location, weights_only=False)
    cfg = TinyChessConfig(**ck["cfg"])
    model = TinyChess(cfg)
    model.load_state_dict(ck["model"])
    if restore_rng and "rng" in ck:
        torch.set_rng_state(ck["rng"]["torch"])
        __import__("random").setstate(ck["rng"]["python"])
        __import__("numpy").random.set_state(ck["rng"]["numpy"])
    return model, cfg, ck


def load_model_only(path, map_location="cpu"):
    ck = torch.load(path, map_location=map_location, weights_only=False)
    cfg = TinyChessConfig(**ck["cfg"])
    model = TinyChess(cfg)
    model.load_state_dict(ck["model"])
    model.eval()
    return model, cfg