// Smoke test : charge le Talker (qwen3 dense + M-RoPE) via le DENSE path engine,
// puis exerce les primitives embeds (prefillEmbeds/decodeEmbed) avec un input bidon.
//
// But : prouver que (a) Patch 2 active M-RoPE (log "mrope sections = [24,20,20,0]"),
// (b) Patch 1 fait tourner llama_decode embeds-only + ressort hidden state + logits.
// Pas de vérification de qualité ici — c'est le job du câblage CP+Decoder (Patch 3).
//
// Usage : ./test_talker <gguf>
//
// Build : copie dans dist/jni/, link avec libkazeia_engine.so dépendances.
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <vector>
#include "llama.h"
#include "ggml-backend.h"

static ggml_backend_dev_t find_htp() {
    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
        auto d = ggml_backend_dev_get(i);
        if (!strcmp(ggml_backend_dev_name(d), "HTP0")) return d;
    }
    return nullptr;
}

int main(int argc, char** argv) {
    if (argc < 2) { printf("usage: %s <gguf>\n", argv[0]); return 1; }

    // dense path : HMX off (cf. kazeia_engine_jni)
    setenv("GGML_HEXAGON_USE_HMX", "0", 1);
    llama_backend_init();

    auto mp = llama_model_default_params();
    // 2e argument optionnel : "cpu" force le path CPU même si HTP dispo (utile pour itérer
    // sans payer l'upload HTP de 1.66GB f32 à chaque smoke test).
    const bool force_cpu = (argc >= 3 && !strcmp(argv[2], "cpu"));
    ggml_backend_dev_t devs[2] = { force_cpu ? nullptr : find_htp(), nullptr };
    if (devs[0]) { mp.n_gpu_layers = 99; mp.devices = devs; printf("HTP0 -> ngl=99\n"); }
    else         { mp.n_gpu_layers = 0;  printf("CPU only%s\n", force_cpu ? " (force)" : " (pas de HTP)"); }

    auto m = llama_model_load_from_file(argv[1], mp);
    if (!m) { printf("model load FAILED\n"); return 1; }
    printf("model loaded. n_embd=%d, n_vocab=%d\n",
           llama_model_n_embd(m), llama_vocab_n_tokens(llama_model_get_vocab(m)));

    auto cp = llama_context_default_params();
    cp.n_ctx = 1024; cp.n_batch = 1024; cp.n_threads = 4;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    cp.embeddings = true;  // équivaut à llama_set_embeddings(ctx, true)

    auto ctx = llama_init_from_model(m, cp);
    if (!ctx) { printf("ctx init FAILED\n"); return 1; }

    const int n_embd  = llama_model_n_embd(m);
    const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(m));

    // --- prefill : 4 positions d'embeds aléatoires (juste pour valider le pipeline,
    //               PAS pour produire du français sensé) ---
    const int T = 4;
    std::vector<float> embd(T * n_embd);
    srand(42);
    for (auto& x : embd) x = (rand() / (float)RAND_MAX - 0.5f) * 0.01f;

    std::vector<llama_pos>     pos(T);
    std::vector<int32_t>       nsd(T, 1);
    std::vector<llama_seq_id>  sid0(T, 0);
    std::vector<llama_seq_id*> sids(T);
    std::vector<int8_t>        lg(T, 0);
    for (int i = 0; i < T; ++i) { pos[i] = i; sids[i] = &sid0[i]; }
    lg[T-1] = 1;

    llama_batch b{};
    b.n_tokens = T; b.token = nullptr; b.embd = embd.data();
    b.pos = pos.data(); b.n_seq_id = nsd.data(); b.seq_id = sids.data(); b.logits = lg.data();

    if (llama_decode(ctx, b) != 0) { printf("prefill embeds FAILED\n"); return 1; }
    printf("prefill OK (T=%d)\n", T);

    const float* h_prefill = llama_get_embeddings_ith(ctx, -1);
    if (!h_prefill) { printf("get_embeddings_ith FAILED\n"); return 1; }
    float h_sum = 0, h_min = h_prefill[0], h_max = h_prefill[0];
    for (int i = 0; i < n_embd; ++i) {
        h_sum += h_prefill[i];
        if (h_prefill[i] < h_min) h_min = h_prefill[i];
        if (h_prefill[i] > h_max) h_max = h_prefill[i];
    }
    printf("hidden[prefill] : mean=%.6f min=%.6f max=%.6f\n", h_sum / n_embd, h_min, h_max);

    const float* logits = llama_get_logits_ith(ctx, -1);
    if (!logits) { printf("get_logits_ith FAILED\n"); return 1; }
    int argmax = 0; float lmax = logits[0];
    for (int i = 1; i < n_vocab; ++i) if (logits[i] > lmax) { lmax = logits[i]; argmax = i; }
    printf("logits[prefill] : argmax=%d (val=%.3f)\n", argmax, lmax);

    // --- decode step : un seul embed de plus ---
    std::vector<float> embd1(n_embd);
    for (auto& x : embd1) x = (rand() / (float)RAND_MAX - 0.5f) * 0.01f;

    llama_pos     p1 = T;
    int32_t       n1 = 1;
    llama_seq_id  s1 = 0; llama_seq_id* sp1 = &s1;
    int8_t        l1 = 1;
    llama_batch b1{};
    b1.n_tokens = 1; b1.token = nullptr; b1.embd = embd1.data();
    b1.pos = &p1; b1.n_seq_id = &n1; b1.seq_id = &sp1; b1.logits = &l1;

    if (llama_decode(ctx, b1) != 0) { printf("decode embeds FAILED\n"); return 1; }
    const float* h_step = llama_get_embeddings_ith(ctx, -1);
    const float* l_step = llama_get_logits_ith(ctx, -1);
    int argmax_s = 0; float lmax_s = l_step[0];
    for (int i = 1; i < n_vocab; ++i) if (l_step[i] > lmax_s) { lmax_s = l_step[i]; argmax_s = i; }
    printf("step OK : argmax=%d (val=%.3f), hidden[0..2]=%.4f %.4f %.4f\n",
           argmax_s, lmax_s, h_step[0], h_step[1], h_step[2]);

    llama_free(ctx);
    llama_model_free(m);
    return 0;
}