// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "core/autograd.h"
#include "train/optimizer.h"

#include <cstdint>
#include <string>
#include <utility>
#include <vector>

namespace forge::train {

// Binary checkpoint: model params (by name), optimizer moments + t, step,
// RNG state, and the config JSON — everything needed for exact resume
// (CLAUDE.md training-loop requirements). Format:
//   magic "FRGE" u32 | version u32 | step i64 | opt_t i64 | rng u64
//   | config_json (u32 len + bytes)
//   | n_params u32 | per param: name (u32+bytes), ndim u32, dims i64[],
//     f32 data
//   | n_opt u32 | per tracked param: m f32[], v f32[]   (0 if never stepped)
struct CheckpointData {
    int64_t step = 0;
    uint64_t rng_state = 0;
    std::string config_json;
};

void save_checkpoint(const std::string& path,
                     const std::vector<std::pair<std::string, Var>>& named_params,
                     AdamW* opt, const CheckpointData& meta);

// Loads params by NAME into an already-constructed model (shape-checked).
// opt may be null (inference). Returns the stored metadata.
CheckpointData load_checkpoint(const std::string& path,
                               const std::vector<std::pair<std::string, Var>>& named_params,
                               AdamW* opt);

// Reads just the config JSON (to construct the model before loading).
std::string read_checkpoint_config(const std::string& path);

} // namespace forge::train