// kazeia_engine_jni.cpp — bridge LLM Kazeia-Engine (llama.cpp upstream + Hexagon).
// Speaker+Thinker, GGUF, prefill NPU (ngl99/HTP0) + decode CPU (8 threads).
// 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 = 99;      // prefill HTP
    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 = 8; cp.n_batch = 512;
    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 prompt, jint maxTok) {
    auto* k = (KEngine*)h; const char* pr = e->GetStringUTFChars(prompt, 0);
    int n = -llama_tokenize(k->v, pr, strlen(pr), nullptr, 0, true, true);
    std::vector<llama_token> t(n); llama_tokenize(k->v, pr, strlen(pr), t.data(), n, true, true);
    e->ReleaseStringUTFChars(prompt, pr);
    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; }