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

#include "core/tensor.h"

#include <cstdint>
#include <random>
#include <string>

namespace forge::train {

// mmap'd uint16 token stream (llm.c format: 256 int32 header {20240520, 1,
// num_tokens}; headerless nanoGPT-style .bin also accepted). Batches are
// random contiguous windows of context_length+1; ids/targets are written
// into caller tensors as i32 (what the GPU kernels take).
class DataLoader {
public:
    DataLoader(const std::string& bin_path, int64_t context_length, uint64_t seed);
    ~DataLoader();

    DataLoader(const DataLoader&) = delete;
    DataLoader& operator=(const DataLoader&) = delete;

    // ids: [B, T] i32, targets: [B*T] i32 (targets = ids shifted by one)
    void next_batch(Tensor& ids, Tensor& targets);
    // Deterministic sequential window (evaluation); wraps around.
    void seq_batch(int64_t index, Tensor& ids, Tensor& targets) const;

    int64_t num_tokens() const { return num_tokens_; }

private:
    void fill(int64_t start, int64_t T, int32_t* ids_row, int32_t* tgt_row) const;

    const uint16_t* tokens_ = nullptr; // into the mmap
    void* map_ = nullptr;
    size_t map_len_ = 0;
    int64_t num_tokens_ = 0;
    int64_t context_ = 0;
    std::mt19937_64 rng_;
};

} // namespace forge::train