// Dump per-token final-norm hidden states from a text-only qwen35 GGUF.
//
// llama.cpp b11407 and later already export these two C++ functions from
// src/llama-ext.h (final RMSNorm / t_h_nextn). They are not declared in the
// installed C header llama.h, and llama-ext.h is not installed with the dylib.
// The declarations below match the exported symbols. No source patch is required
// on b11407+. An older tree that lacks the symbols needs a newer llama.cpp,
// not a patch against this file.
//
//   llama_set_embeddings_nextn(ctx, true, masked=false)
//     masked=false: one row per token in the batch, ignoring batch.logits.
//     embeddings stays false. batch.logits is 1 only on the last token of each
//     ubatch so the unused LM head does not reserve n_vocab rows per token.
//
// tokens.bin: int32 nrec, then (int32 ntok, int32 ids[ntok]) * nrec
// stdout:     (int32 ntok, int32 n_embd, float h[ntok * n_embd]) * nrec
//
//   c++ -std=c++17 -O2 \
//     -I llama.cpp/include -I llama.cpp/ggml/include \
//     hsdump.cpp -o hsdump \
//     -L /path/to/llama.cpp/build/bin -lllama \
//     -Wl,-rpath,/path/to/llama.cpp/build/bin
//   hsdump model.gguf tokens.bin [--ngl N] [--ctx N] [--ubatch N]
#include "llama.h"

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>

// llama-ext.h is C++ and not installed with the dylib. These two match b11407's exported symbols.
void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked);
float * llama_get_embeddings_nextn(struct llama_context * ctx);

static void read_exact(FILE * f, void * dst, size_t n, const char * what) {
    auto * p = (uint8_t *) dst;
    size_t got = 0;
    while (got < n) {
        size_t r = std::fread(p + got, 1, n - got, f);
        if (r == 0) {
            std::fprintf(stderr, "short read (%s)\n", what);
            std::exit(2);
        }
        got += r;
    }
}

static void write_exact(const void * src, size_t n) {
    auto * p = (const uint8_t *) src;
    size_t put = 0;
    while (put < n) {
        size_t w = std::fwrite(p + put, 1, n - put, stdout);
        if (w == 0) {
            std::fprintf(stderr, "short stdout\n");
            std::exit(2);
        }
        put += w;
    }
}

struct Rec {
    std::vector<llama_token> ids;
};

static std::vector<Rec> load_tokens(const char * path, int32_t & max_len) {
    FILE * f = std::fopen(path, "rb");
    if (!f) {
        std::perror(path);
        std::exit(2);
    }
    int32_t nrec = 0;
    read_exact(f, &nrec, 4, "nrec");
    if (nrec < 0 || nrec > 1000000) {
        std::fprintf(stderr, "bad nrec %d\n", nrec);
        std::exit(2);
    }
    std::vector<Rec> recs(nrec);
    max_len = 0;
    for (int32_t i = 0; i < nrec; ++i) {
        int32_t n = 0;
        read_exact(f, &n, 4, "ntok");
        if (n <= 0 || n > 1000000) {
            std::fprintf(stderr, "bad ntok %d at record %d\n", n, i);
            std::exit(2);
        }
        recs[i].ids.resize(n);
        read_exact(f, recs[i].ids.data(), (size_t) n * sizeof(llama_token), "ids");
        if (n > max_len) max_len = n;
    }
    std::fclose(f);
    return recs;
}

int main(int argc, char ** argv) {
    if (argc < 3) {
        std::fprintf(stderr, "usage: hsdump model.gguf tokens.bin [--ngl N] [--ctx N] [--ubatch N]\n");
        return 2;
    }
    const char * model = argv[1];
    const char * tokens_path = argv[2];
    int ngl = 99;
    int ctx_arg = 0;
    int ubatch = 512;
    for (int i = 3; i < argc; ++i) {
        if (!std::strcmp(argv[i], "--ngl") && i + 1 < argc) ngl = std::atoi(argv[++i]);
        else if (!std::strcmp(argv[i], "--ctx") && i + 1 < argc) ctx_arg = std::atoi(argv[++i]);
        else if (!std::strcmp(argv[i], "--ubatch") && i + 1 < argc) ubatch = std::atoi(argv[++i]);
        else {
            std::fprintf(stderr, "unknown arg %s\n", argv[i]);
            return 2;
        }
    }

    int32_t max_len = 0;
    std::vector<Rec> recs = load_tokens(tokens_path, max_len);
    const int n_ctx = ctx_arg > 0 ? ctx_arg : max_len + 8;
    if (max_len > n_ctx) {
        std::fprintf(stderr, "record length %d exceeds n_ctx %d\n", max_len, n_ctx);
        return 2;
    }
    std::fprintf(stderr, "records %zu max_len %d n_ctx %d ubatch %d ngl %d\n",
                 recs.size(), max_len, n_ctx, ubatch, ngl);

    llama_backend_init();
    llama_model_params mp = llama_model_default_params();
    mp.n_gpu_layers = ngl;
    llama_model * mdl = llama_model_load_from_file(model, mp);
    if (!mdl) {
        std::fprintf(stderr, "failed to load %s\n", model);
        return 1;
    }
    llama_context_params cp = llama_context_default_params();
    cp.n_ctx = (uint32_t) n_ctx;
    cp.n_batch = (uint32_t) ubatch;
    cp.n_ubatch = (uint32_t) ubatch;
    cp.n_seq_max = 1;
    cp.embeddings = false;
    cp.pooling_type = LLAMA_POOLING_TYPE_NONE;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    llama_context * ctx = llama_init_from_model(mdl, cp);
    if (!ctx) {
        std::fprintf(stderr, "failed to init context\n");
        return 1;
    }
    llama_set_embeddings_nextn(ctx, true, /*masked=*/false);
    const int n_embd = llama_model_n_embd(mdl);
    std::fprintf(stderr, "n_embd %d\n", n_embd);

    llama_batch batch = llama_batch_init(ubatch, 0, 1);
    std::vector<float> hidden;
    for (size_t ri = 0; ri < recs.size(); ++ri) {
        const auto & ids = recs[ri].ids;
        const int32_t n = (int32_t) ids.size();
        hidden.assign((size_t) n * (size_t) n_embd, 0.f);
        llama_memory_clear(llama_get_memory(ctx), true);
        for (int32_t pos = 0; pos < n; ) {
            const int32_t ntok = (n - pos) < ubatch ? (n - pos) : ubatch;
            batch.n_tokens = ntok;
            for (int32_t i = 0; i < ntok; ++i) {
                batch.token[i] = ids[pos + i];
                batch.pos[i] = pos + i;
                batch.n_seq_id[i] = 1;
                batch.seq_id[i][0] = 0;
                batch.logits[i] = (i == ntok - 1) ? 1 : 0;
            }
            const int rc = llama_decode(ctx, batch);
            if (rc != 0) {
                std::fprintf(stderr, "decode failed record %zu pos %d rc %d\n", ri, pos, rc);
                return 1;
            }
            float * h = llama_get_embeddings_nextn(ctx);
            if (!h) {
                std::fprintf(stderr, "null nextn embeddings record %zu pos %d\n", ri, pos);
                return 1;
            }
            std::memcpy(hidden.data() + (size_t) pos * (size_t) n_embd, h,
                        (size_t) ntok * (size_t) n_embd * sizeof(float));
            pos += ntok;
        }
        int32_t hdr[2] = { n, n_embd };
        write_exact(hdr, sizeof(hdr));
        write_exact(hidden.data(), hidden.size() * sizeof(float));
        std::fflush(stdout);
        if ((ri + 1) % 10 == 0 || ri + 1 == recs.size()) {
            std::fprintf(stderr, "dumped %zu/%zu\n", ri + 1, recs.size());
        }
    }
    llama_batch_free(batch);
    llama_free(ctx);
    llama_model_free(mdl);
    llama_backend_free();
    std::fprintf(stderr, "HSDUMP_OK\n");
    return 0;
}