#ifndef XENITH_CONVERTER_H
#define XENITH_CONVERTER_H
#include <cstdint>
#include <fstream>
#include <ios>
#include <string>
#include <unordered_map>
#include <vector>

enum cache_type
{
    K_cache,
    V_cache
};

const uint32_t BFR_MAGIC = 0x31524642;

struct ModelConfig {
    uint32_t magic = 0x31524642;
    uint32_t vocab_size;
    uint32_t embedding_dim;
    uint32_t num_heads;
    uint32_t head_dim;
    uint32_t num_layers;
    uint32_t max_ctx;
};

struct Model {
    ModelConfig config;

    std::unordered_map<std::string, uint32_t> vocab;

    std::vector<float> w_emb;
    std::vector<float> w_qkv;
    std::vector<float> w_o;

    FILE* fd = nullptr;

    uint64_t offset_emb;
    uint64_t offset_qkv;
    uint64_t offset_o;

    float* kv_cache = nullptr;
    int ctx_len = 0;


    void close_fd();

    void load(size_t num, std::vector<float>& target, uint64_t offset);

    void load_edm();
    void load_qkv();
    void load_o();

    void free_edm();
    void free_qkv();
    void free_o();

    void allocate_kv_cache();
    void free_kv_cache();

    float* get_kv(uint64_t token, uint32_t layer, cache_type kv, uint32_t head, uint32_t dim);
};


bool save_to_bfr(const std::string& filename, const Model& model);
Model load_from_bfr(const std::string& filename);

#endif //XENITH_CONVERTER_H