// kazeia_engine_jni.cpp — bridge LLM Kazeia-Engine (llama.cpp upstream + Hexagon).
// Speaker+Thinker, GGUF, prefill NPU (ngl99/HTP0) + decode CPU (t4+fa = 14t/s; t8=contention).
// thinking off via template applique cote Kotlin (/no_think) OU enable_thinking=false.
// STT reste ORT-QAIRT. Build: link libllama+libcommon+ggml*; jniLibs htp auto-V79.
#include <jni.h>
#include <string>
#include <vector>
#include "llama.h"

struct KEngine { llama_model* m; llama_context* c; const llama_vocab* v; llama_sampler* s; };

extern "C" JNIEXPORT jlong JNICALL
Java_com_kazeia_llm_EngineJni_load(JNIEnv* e, jobject, jstring path, jint nctx) {
    const char* p = e->GetStringUTFChars(path, 0);
    llama_backend_init();
    auto mp = llama_model_default_params(); mp.n_gpu_layers = 0; // decode CPU pur 14t/s (ngl99=ping-pong 0.2)
    auto m = llama_model_load_from_file(p, mp); e->ReleaseStringUTFChars(path, p);
    if (!m) return 0;
    auto cp = llama_context_default_params(); cp.n_ctx = nctx; cp.n_threads = 4; cp.n_batch = 512;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; // t4+fa = 14t/s decode (t8=9.8 contention; fa coupe trafic KV)
    cp.type_k = GGML_TYPE_Q8_0; cp.type_v = GGML_TYPE_Q8_0; // KV q8: 14->15.5 dense (>pte), 8.5->9.8 hyb, quasi-lossless
    auto* k = new KEngine{m, llama_init_from_model(m, cp), llama_model_get_vocab(m),
                          llama_sampler_init_greedy()};
    return (jlong)k;
}

extern "C" JNIEXPORT jstring JNICALL
Java_com_kazeia_llm_EngineJni_generate(JNIEnv* e, jobject, jlong h, jstring sys, jstring usr, jint maxTok) {
    auto* k = (KEngine*)h;
    const char* sp = e->GetStringUTFChars(sys, 0); const char* up = e->GetStringUTFChars(usr, 0);
    // ChatML Qwen3.5 + <think></think> vide = thinking OFF deterministe (sinon boucle infinie)
    std::string p = "<|im_start|>system\n"; p+=sp; p+="<|im_end|>\n<|im_start|>user\n"; p+=up;
    p += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
    e->ReleaseStringUTFChars(sys, sp); e->ReleaseStringUTFChars(usr, up);
    int n = -llama_tokenize(k->v, p.c_str(), p.size(), nullptr, 0, true, true);
    std::vector<llama_token> t(n); llama_tokenize(k->v, p.c_str(), p.size(), t.data(), n, true, true);
    llama_batch b = llama_batch_get_one(t.data(), n);            // prefill batch -> HTP
    std::string out; char buf[256];
    for (int i = 0; i < maxTok && llama_decode(k->c, b) == 0; ++i) {
        llama_token id = llama_sampler_sample(k->s, k->c, -1);   // decode -> CPU
        if (llama_vocab_is_eog(k->v, id)) break;
        int l = llama_token_to_piece(k->v, id, buf, sizeof buf, 0, true);
        if (l > 0) out.append(buf, l);
        b = llama_batch_get_one(&id, 1);
    }
    return e->NewStringUTF(out.c_str());
}

extern "C" JNIEXPORT void JNICALL
Java_com_kazeia_llm_EngineJni_reset(JNIEnv*, jobject, jlong h){ llama_memory_clear(llama_get_memory(((KEngine*)h)->c), true); }
extern "C" JNIEXPORT void JNICALL
Java_com_kazeia_llm_EngineJni_free(JNIEnv*, jobject, jlong h){ auto*k=(KEngine*)h; llama_sampler_free(k->s); llama_free(k->c); llama_model_free(k->m); delete k; }
