// sh-states: dump the hidden states a Standard One schema head reads (final norm output + one decoder layer output)
// from a llama.cpp GGUF model, for already-tokenized inputs.
//
//   sh-states -m model.gguf --tap 25 [-t threads] [-c n_ctx] [-b chunk]
//
// Protocol (line based, stdin -> stdout): each request line is
//   <ids.bin> <out_prefix>
// where ids.bin holds little-endian int32 token ids (BOS included, exactly as the renderer produced them). For every
// request the program writes <out_prefix>.final.f32 and <out_prefix>.tap.f32 (row-major float32 [n_tokens, n_embd])
// and answers "ok <n_tokens> <n_embd>" (or "err <message>"). The KV cache is cleared between requests.
//
// The final state is llama.cpp's per-token embedding output (embeddings=true, pooling NONE: the final RMSNorm output =
// HF last_hidden_state). The tap is taken with the scheduler eval callback on "l_out-<tap>" (= the output of decoder
// layer <tap>, as a forward hook on layers[tap] sees it). The prompt is decoded in chunks with every token marked as
// an output, so both cover all tokens.
#include "llama.h"
#include "ggml.h"
#include "ggml-backend.h"

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

struct capture {
    std::string tap_name;
    std::vector<float> final_rows, tap_rows;
    int64_t n_embd = 0;
};

static bool eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
    auto * cap = (capture *) user_data;
    const bool want = cap->tap_name == t->name;
    if (ask) {
        return want;
    }
    if (!want) {
        return true;
    }
    if (t->type != GGML_TYPE_F32 || ggml_n_dims(t) > 2) {
        fprintf(stderr, "unexpected tensor %s type %d dims %d\n", t->name, (int) t->type, ggml_n_dims(t));
        return false;
    }
    const int64_t n_embd = t->ne[0], n_tok = t->ne[1];
    cap->n_embd = n_embd;
    std::vector<float> buf(n_embd * n_tok);
    if (t->nb[1] == (size_t) n_embd * sizeof(float)) {
        ggml_backend_tensor_get(t, buf.data(), 0, buf.size() * sizeof(float));
    } else {   // non-contiguous rows
        for (int64_t r = 0; r < n_tok; r++) {
            ggml_backend_tensor_get(t, buf.data() + r * n_embd, r * t->nb[1], n_embd * sizeof(float));
        }
    }
    cap->tap_rows.insert(cap->tap_rows.end(), buf.begin(), buf.end());
    return true;
}

static bool write_f32(const std::string & path, const std::vector<float> & v) {
    std::ofstream f(path, std::ios::binary);
    f.write((const char *) v.data(), v.size() * sizeof(float));
    return (bool) f;
}

int main(int argc, char ** argv) {
    std::string model_path;
    int tap = 25, threads = 8, n_ctx = 34000, chunk = 512;
    for (int i = 1; i < argc; i++) {
        std::string a = argv[i];
        auto next = [&]() { return std::string(i + 1 < argc ? argv[++i] : ""); };
        if (a == "-m") model_path = next();
        else if (a == "--tap") tap = std::stoi(next());
        else if (a == "-t") threads = std::stoi(next());
        else if (a == "-c") n_ctx = std::stoi(next());
        else if (a == "-b") chunk = std::stoi(next());
        else { fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; }
    }
    if (model_path.empty()) { fprintf(stderr, "usage: sh-states -m model.gguf --tap 25\n"); return 2; }

    llama_backend_init();
    llama_model_params mp = llama_model_default_params();
    llama_model * model = llama_model_load_from_file(model_path.c_str(), mp);
    if (!model) { fprintf(stderr, "failed to load %s\n", model_path.c_str()); return 1; }
    if (tap < 0 || tap >= llama_model_n_layer(model)) { fprintf(stderr, "tap out of range\n"); return 2; }

    capture cap;
    cap.tap_name = "l_out-" + std::to_string(tap);
    llama_context_params cp = llama_context_default_params();
    cp.n_ctx = n_ctx;
    cp.n_batch = chunk;
    cp.n_ubatch = chunk;
    cp.n_threads = threads;
    cp.n_threads_batch = threads;
    cp.embeddings = true;
    cp.pooling_type = LLAMA_POOLING_TYPE_NONE;
    cp.cb_eval = eval_cb;
    cp.cb_eval_user_data = &cap;
    llama_context * ctx = llama_init_from_model(model, cp);
    if (!ctx) { fprintf(stderr, "failed to create context\n"); return 1; }
    printf("ready n_layer=%d n_embd=%d tap=%d\n", llama_model_n_layer(model), llama_model_n_embd(model), tap);
    fflush(stdout);

    std::string line;
    while (std::getline(std::cin, line)) {
        std::istringstream ss(line);
        std::string ids_path, out_prefix;
        if (!(ss >> ids_path >> out_prefix)) { printf("err bad request line\n"); fflush(stdout); continue; }
        std::ifstream f(ids_path, std::ios::binary | std::ios::ate);
        if (!f) { printf("err cannot read %s\n", ids_path.c_str()); fflush(stdout); continue; }
        const size_t bytes = f.tellg();
        f.seekg(0);
        std::vector<llama_token> ids(bytes / sizeof(int32_t));
        f.read((char *) ids.data(), bytes);
        const int n = (int) ids.size();
        if (n == 0 || n > n_ctx) { printf("err %d tokens (context %d)\n", n, n_ctx); fflush(stdout); continue; }
        llama_memory_clear(llama_get_memory(ctx), true);
        cap.final_rows.clear();
        cap.tap_rows.clear();
        bool ok = true;
        llama_batch batch = llama_batch_init(chunk, 0, 1);
        for (int s = 0; s < n && ok; s += chunk) {
            const int m = std::min(chunk, n - s);
            batch.n_tokens = m;
            for (int j = 0; j < m; j++) {
                batch.token[j] = ids[s + j];
                batch.pos[j] = s + j;
                batch.n_seq_id[j] = 1;
                batch.seq_id[j][0] = 0;
                batch.logits[j] = 1;
            }
            if (llama_decode(ctx, batch) != 0) { ok = false; break; }
            const float * emb = llama_get_embeddings(ctx);
            if (!emb) { ok = false; break; }
            const int64_t h = llama_model_n_embd(model);
            cap.final_rows.insert(cap.final_rows.end(), emb, emb + (size_t) m * h);
        }
        llama_batch_free(batch);
        const size_t want = (size_t) n * cap.n_embd;
        if (!ok || cap.final_rows.size() != want || cap.tap_rows.size() != want) {
            printf("err decode=%d final=%zu tap=%zu want=%zu\n", ok, cap.final_rows.size(), cap.tap_rows.size(), want);
            fflush(stdout);
            continue;
        }
        if (!write_f32(out_prefix + ".final.f32", cap.final_rows) || !write_f32(out_prefix + ".tap.f32", cap.tap_rows)) {
            printf("err write failed\n");
        } else {
            printf("ok %d %lld\n", n, (long long) cap.n_embd);
        }
        fflush(stdout);
    }
    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();
    return 0;
}