// Offline independent-runtime fixture. Build against the pinned Flint llama.cpp
// branch documented in docs/flint_quantization.md. Input: little-endian uint32
// document count, sequence length, then int32 token IDs. Output: FLINTLOG magic,
// uint32 document count/length/vocabulary, then all FP32 logits, in row order.
#include "llama.h"
#include <cstdint>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <vector>

int main(int argc, char ** argv) {
    if (argc != 4) {
        std::cerr << "usage: flint-llama-fixture model.gguf tokens.bin logits.bin\n";
        return 2;
    }
    std::ifstream input(argv[2], std::ios::binary);
    uint32_t documents = 0, length = 0;
    input.read(reinterpret_cast<char *>(&documents), 4);
    input.read(reinterpret_cast<char *>(&length), 4);
    if (!input || !documents || !length || length > 16384) {
        throw std::runtime_error("invalid fixture header");
    }
    llama_backend_init();
    auto mp = llama_model_default_params();
    mp.n_gpu_layers = 999;
    auto * model = llama_model_load_from_file(argv[1], mp);
    if (!model) throw std::runtime_error("model load failed");
    auto params = llama_context_default_params();
    params.n_ctx = length;
    params.n_batch = length;
    params.n_ubatch = length;
    params.n_seq_max = 1;
    params.n_threads = 8;
    params.n_threads_batch = 8;
    params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
    params.type_k = GGML_TYPE_F32;
    params.type_v = GGML_TYPE_F32;
    auto * ctx = llama_init_from_model(model, params);
    if (!ctx) throw std::runtime_error("context creation failed");
    uint32_t vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    std::ofstream output(argv[3], std::ios::binary);
    output.write("FLINTLOG", 8);
    for (auto value : {documents, length, vocab}) output.write(reinterpret_cast<char *>(&value), 4);
    auto batch = llama_batch_init(length, 0, 1);
    for (uint32_t doc = 0; doc < documents; ++doc) {
        llama_memory_clear(llama_get_memory(ctx), true);
        batch.n_tokens = length;
        input.read(reinterpret_cast<char *>(batch.token), length * sizeof(llama_token));
        if (!input) throw std::runtime_error("short token fixture");
        for (uint32_t i = 0; i < length; ++i) {
            batch.pos[i] = i;
            batch.n_seq_id[i] = 1;
            batch.seq_id[i][0] = 0;
            batch.logits[i] = 1;
        }
        if (llama_decode(ctx, batch) != 0) throw std::runtime_error("decode failed");
        for (uint32_t i = 0; i < length; ++i) {
            auto * logits = llama_get_logits_ith(ctx, i);
            if (!logits) throw std::runtime_error("missing logits");
            output.write(reinterpret_cast<char *>(logits), vocab * sizeof(float));
        }
        if (!output) throw std::runtime_error("logits write failed");
        std::cout << "completed document " << doc << '\n' << std::flush;
    }
    llama_batch_free(batch);
    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();
}