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

#include "nn/config.h"
#include "nn/transformer.h"
#include "train/dataloader.h"
#include "train/optimizer.h"

#include <memory>
#include <string>

namespace forge::train {

// Training loop: gradient accumulation, warmup+cosine LR, global-norm clip,
// periodic eval + resumable checkpoints, stdout + CSV logging (step, loss,
// tokens/sec, lr, grad norm — CLAUDE.md training-loop requirements).
class Trainer {
public:
    Trainer(Config cfg, const std::string& data_dir, const std::string& out_dir,
            const std::string& config_json);

    // resume_from: checkpoint path or empty.
    void train(const std::string& resume_from);

private:
    float eval_loss();
    void save(int64_t step);

    Config cfg_;
    std::string out_dir_;
    std::string config_json_;
    std::unique_ptr<nn::Transformer> model_;
    std::unique_ptr<AdamW> opt_;
    std::unique_ptr<DataLoader> train_data_;
    std::unique_ptr<DataLoader> val_data_;
};

} // namespace forge::train