// pog-extract: Project Omni-Gemma feature extractor
//
// Intercepts per-layer residual-stream hidden states ("l_out-<il>") from a
// frozen QAT GGUF backbone during batched prefill, via the ggml scheduler
// eval callback (same mechanism as llama-imatrix). The LM head is skipped
// entirely by requesting zero outputs from llama_decode.
//
// For each input text it caches, per intercepted layer:
//   slot 0            : mean of instruction-prefix tokens (incl. BOS)
//   slots 1..K        : means of K contiguous content segments
//   slot K+1          : last-token hidden state
// as fp16, record layout [n_layers][K+2][n_embd].
//
// Input : TSV lines  "id\ttype\ttext"  (type: q = query, d = document)
// Output: <out>.bin (fp16 records) + <out>.idx (one line per record: id type n_tok)
//
// Usage: pog-extract <model.gguf> <input.tsv> <out-prefix>
//                    [--layers 28,32,36] [--segments 4] [--max-tokens 320]
//                    [--ubatch 8192] [--ngl 999]
//
// Server mode (model stays loaded; one line per command on stdin):
//   pog-extract <model.gguf> --serve [opts]
//     RUN <input.tsv> <out-prefix> [max_tokens]   -> "OK <n_records>"
//     QUIT

#include "llama.h"
#include "ggml.h"
#include "ggml-backend.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cinttypes>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <string>
#include <vector>

struct capture_ctx {
    std::vector<int>                    layers;     // intercepted layer indices
    std::map<int, std::vector<float>>   data;       // layer -> [n_tokens * n_embd] f32
    int                                 n_embd      = 0;
    int                                 n_tokens    = 0; // tokens in current decode
    bool                                armed       = false;
};

static bool cb_capture(struct ggml_tensor * t, bool ask, void * user_data) {
    auto * cap = (capture_ctx *) user_data;
    if (!cap->armed) {
        return false;
    }
    const char * name = t->name;
    if (strncmp(name, "l_out-", 6) != 0) {
        return false;
    }
    const int il = atoi(name + 6);
    bool wanted = false;
    for (int l : cap->layers) {
        if (l == il) { wanted = true; break; }
    }
    if (!wanted) {
        return false;
    }
    if (ask) {
        return true; // yes, deliver this tensor's data
    }
    if (t->ne[0] != cap->n_embd || t->ne[1] != cap->n_tokens) {
        fprintf(stderr, "FATAL: %s shape [%" PRId64 ",%" PRId64 "] != expected [%d,%d] "
                        "(ubatch split? raise --ubatch)\n",
                name, t->ne[0], t->ne[1], cap->n_embd, cap->n_tokens);
        exit(3);
    }
    if (t->type != GGML_TYPE_F32) {
        fprintf(stderr, "FATAL: %s is not F32 (type=%d)\n", name, (int) t->type);
        exit(3);
    }
    auto & buf = cap->data[il];
    buf.resize((size_t) cap->n_embd * cap->n_tokens);
    ggml_backend_tensor_get(t, buf.data(), 0, buf.size() * sizeof(float));
    return true;
}

struct pending_text {
    std::string id;
    char        type;       // 'q' or 'd'
    int         seq_start;  // token offset within batch
    int         n_tok;      // total tokens (prefix + content)
    int         n_prefix;   // prefix tokens (incl. BOS)
};

int main(int argc, char ** argv) {
    if (argc < 4) {
        fprintf(stderr, "usage: %s <model.gguf> <input.tsv> <out-prefix> [opts]\n", argv[0]);
        return 1;
    }
    const char * model_path = argv[1];
    const bool   serve      = (strcmp(argv[2], "--serve") == 0);
    const char * input_path = serve ? nullptr : argv[2];
    std::string  out_prefix = serve ? ""      : argv[3];

    std::vector<int> layers = {28, 32, 36};
    int n_segments  = 4;
    int max_tokens  = 320;
    int n_ubatch    = 8192;
    int n_gpu_layers = 999;

    for (int i = serve ? 3 : 4; i < argc; i++) {
        std::string a = argv[i];
        if (a == "--layers" && i + 1 < argc) {
            layers.clear();
            std::stringstream ss(argv[++i]);
            std::string item;
            while (std::getline(ss, item, ',')) layers.push_back(atoi(item.c_str()));
        } else if (a == "--segments"   && i + 1 < argc) { n_segments  = atoi(argv[++i]); }
        else if (a == "--max-tokens"   && i + 1 < argc) { max_tokens  = atoi(argv[++i]); }
        else if (a == "--ubatch"       && i + 1 < argc) { n_ubatch    = atoi(argv[++i]); }
        else if (a == "--ngl"          && i + 1 < argc) { n_gpu_layers = atoi(argv[++i]); }
        else { fprintf(stderr, "unknown arg %s\n", a.c_str()); return 1; }
    }

    llama_backend_init();

    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers = n_gpu_layers;
    llama_model * model = llama_model_load_from_file(model_path, mparams);
    if (!model) { fprintf(stderr, "failed to load model\n"); return 1; }

    const llama_vocab * vocab  = llama_model_get_vocab(model);
    const int           n_embd = llama_model_n_embd(model);

    capture_ctx cap;
    cap.layers = layers;
    cap.n_embd = n_embd;

    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx             = n_ubatch;
    cparams.n_batch           = n_ubatch;
    cparams.n_ubatch          = n_ubatch;
    cparams.n_seq_max         = 256;
    cparams.embeddings        = false;
    cparams.kv_unified        = true;  // keep split_simple: ubatch token order == insertion order
    cparams.cb_eval           = cb_capture;
    cparams.cb_eval_user_data = &cap;
    cparams.no_perf           = true;

    llama_context * ctx = llama_init_from_model(model, cparams);
    if (!ctx) { fprintf(stderr, "failed to create context\n"); return 1; }
    llama_memory_t mem = llama_get_memory(ctx);

    // tokenize the two instruction prefixes once
    auto tokenize = [&](const std::string & s, bool add_bos) {
        std::vector<llama_token> toks(s.size() + 16);
        int n = llama_tokenize(vocab, s.c_str(), (int32_t) s.size(),
                               toks.data(), (int32_t) toks.size(), add_bos, false);
        if (n < 0) { toks.resize(-n);
            n = llama_tokenize(vocab, s.c_str(), (int32_t) s.size(),
                               toks.data(), (int32_t) toks.size(), add_bos, false); }
        toks.resize(n);
        return toks;
    };
    const std::string prefix_q = "[Task: Retrieval] [Type: Query]\n";
    const std::string prefix_d = "[Task: Retrieval] [Type: Document]\n";
    const auto ptoks_q = tokenize(prefix_q, true);
    const auto ptoks_d = tokenize(prefix_d, true);

    const int n_slots = n_segments + 2;
    const size_t rec_floats = (size_t) layers.size() * n_slots * n_embd;
    std::vector<ggml_fp16_t> rec(rec_floats);

    llama_batch batch = llama_batch_init(n_ubatch, 0, 256);

    auto run_file = [&](const std::string & in_path, const std::string & out_pref,
                        int max_tok) -> long {
    std::ifstream fin(in_path);
    if (!fin) { fprintf(stderr, "cannot open %s\n", in_path.c_str()); return -1; }
    std::ofstream fbin(out_pref + ".bin", std::ios::binary);
    std::ofstream fidx(out_pref + ".idx");

    std::vector<pending_text>             pend;
    std::vector<std::vector<llama_token>> pend_toks;
    int  batch_tokens = 0;
    long n_done = 0, n_tok_total = 0;
    const int64_t t_start = ggml_time_us();

    auto flush = [&]() {
        if (pend.empty()) return;
        batch.n_tokens = 0;
        for (size_t s = 0; s < pend.size(); s++) {
            const auto & toks = pend_toks[s];
            for (int j = 0; j < (int) toks.size(); j++) {
                const int i = batch.n_tokens++;
                batch.token[i]    = toks[j];
                batch.pos[i]      = j;
                batch.n_seq_id[i] = 1;
                batch.seq_id[i][0] = (llama_seq_id) s;
                batch.logits[i]   = 0; // no outputs anywhere: LM head is skipped
            }
        }
        cap.n_tokens = batch.n_tokens;
        cap.data.clear();
        cap.armed = true;
        if (llama_decode(ctx, batch) != 0) {
            fprintf(stderr, "FATAL: llama_decode failed\n");
            exit(2);
        }
        cap.armed = false;
        for (int l : layers) {
            if (cap.data.find(l) == cap.data.end()) {
                fprintf(stderr, "FATAL: layer %d not captured\n", l);
                exit(3);
            }
        }
        // pool + write
        for (const auto & p : pend) {
            size_t w = 0;
            for (int l : layers) {
                const float * d = cap.data[l].data();
                auto emit_mean = [&](int a, int b) { // token positions [a,b) within text
                    if (b <= a) b = a + 1;           // degenerate: fall back to single token
                    std::vector<double> acc(n_embd, 0.0);
                    for (int t = a; t < b; t++) {
                        const float * v = d + (size_t) (p.seq_start + t) * n_embd;
                        for (int e = 0; e < n_embd; e++) acc[e] += v[e];
                    }
                    const double inv = 1.0 / (b - a);
                    for (int e = 0; e < n_embd; e++)
                        rec[w++] = ggml_fp32_to_fp16((float) (acc[e] * inv));
                };
                emit_mean(0, p.n_prefix);                       // prefix mean
                const int c0 = p.n_prefix, cn = p.n_tok - p.n_prefix;
                for (int s = 0; s < n_segments; s++) {           // content segment means
                    const int a = c0 + (int) ((int64_t) cn * s       / n_segments);
                    const int b = c0 + (int) ((int64_t) cn * (s + 1) / n_segments);
                    emit_mean(a, b);
                }
                emit_mean(p.n_tok - 1, p.n_tok);                 // last token
            }
            fbin.write((const char *) rec.data(), rec.size() * sizeof(ggml_fp16_t));
            fidx << p.id << '\t' << p.type << '\t' << p.n_tok << '\n';
        }
        n_done += (long) pend.size();
        n_tok_total += batch.n_tokens;
        llama_memory_clear(mem, true);
        pend.clear();
        pend_toks.clear();
        batch_tokens = 0;
        if (n_done % 4096 < (long) cparams.n_seq_max) {
            const double dt = (ggml_time_us() - t_start) / 1e6;
            fprintf(stderr, "[pog] %ld texts | %.0f tok/s | %.1f min\n",
                    n_done, n_tok_total / dt, dt / 60.0);
        }
    };

    std::string line;
    while (std::getline(fin, line)) {
        if (line.empty()) continue;
        const size_t t1 = line.find('\t');
        const size_t t2 = line.find('\t', t1 + 1);
        if (t1 == std::string::npos || t2 == std::string::npos) continue;
        const std::string id   = line.substr(0, t1);
        const char        type = line[t1 + 1];
        const std::string text = line.substr(t2 + 1);

        const auto & ptoks = (type == 'q') ? ptoks_q : ptoks_d;
        auto ctoks = tokenize(text, false);
        const int max_content = max_tok - (int) ptoks.size();
        if ((int) ctoks.size() > max_content) ctoks.resize(max_content);
        if (ctoks.empty()) ctoks.push_back(llama_vocab_nl(vocab));

        std::vector<llama_token> toks(ptoks);
        toks.insert(toks.end(), ctoks.begin(), ctoks.end());

        if (batch_tokens + (int) toks.size() > n_ubatch || (int) pend.size() >= 255) {
            flush();
        }
        pending_text p;
        p.id = id; p.type = type;
        p.n_tok = (int) toks.size();
        p.n_prefix = (int) ptoks.size();
        p.seq_start = batch_tokens;
        batch_tokens += (int) toks.size();
        pend.push_back(p);
        pend_toks.push_back(std::move(toks));
    }
    flush();

    const double dt = (ggml_time_us() - t_start) / 1e6;
    fprintf(stderr, "[pog] DONE %ld texts, %ld tokens, %.0f tok/s, %.1f min, "
                    "record = %zu fp16 (%zu layers x %d slots x %d dim)\n",
            n_done, n_tok_total, n_tok_total / dt, dt / 60.0,
            rec_floats, layers.size(), n_slots, n_embd);
    return n_done;
    };

    if (serve) {
        fprintf(stdout, "READY\n");
        fflush(stdout);
        std::string cmd;
        while (std::getline(std::cin, cmd)) {
            if (cmd == "QUIT" || cmd.empty()) break;
            std::stringstream ss(cmd);
            std::string verb, in_path, out_pref;
            int max_tok = max_tokens;
            ss >> verb >> in_path >> out_pref;
            if (!(ss >> max_tok)) max_tok = max_tokens;
            if (verb != "RUN" || in_path.empty() || out_pref.empty()) {
                fprintf(stdout, "ERR bad command\n"); fflush(stdout);
                continue;
            }
            const long got = run_file(in_path, out_pref, max_tok);
            fprintf(stdout, got >= 0 ? "OK %ld\n" : "ERR run failed\n", got);
            fflush(stdout);
        }
    } else {
        if (run_file(input_path, out_prefix, max_tokens) < 0) return 1;
    }

    llama_batch_free(batch);
    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();
    return 0;
}