// Summarize-LLM thin JNI layer over llama.cpp (milestone 1: post-meeting minutes).
//
// This is the ONLY native code added for the summarize feature. It is compiled
// together with a pinned llama.cpp checkout into a single prebuilt shared
// library (libsummarize-llm-jni.so) so that the app's Gradle build keeps using
// only prebuilt .so files (no externalNativeBuild, no NDK required on the
// developer machine). See tools/llama-jni/build_arm64.sh.
//
// Safety rules (mirrors the project's "missing asset => native exit()" lesson):
//   - Every entry point validates inputs and throws a Java exception instead
//     of crashing; model loading never calls exit().
//   - The caller (Kotlin) MUST check that the GGUF file exists before
//     nativeNewFromFile(); we double-check with access() here as well.

#include <jni.h>
#include <android/log.h>

#include <atomic>
#include <mutex>
#include <string>
#include <vector>
#include <algorithm>
#include <cmath>
#include <random>
#include <unistd.h>

#include "llama.h"

#define LOG_TAG "SummarizeLlmJni"
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, LOG_TAG, __VA_ARGS__)

namespace {

struct Engine {
    llama_model * model = nullptr;
    llama_context * ctx = nullptr;
    const llama_vocab * vocab = nullptr;
    std::mutex mu;          // generate() is single-flight
    std::atomic<bool> abort{false};
};

void throwJava(JNIEnv * env, const char * msg) {
    jclass cls = env->FindClass("java/lang/RuntimeException");
    if (cls != nullptr) {
        env->ThrowNew(cls, msg);
    }
}

std::string jstringToUtf8(JNIEnv * env, jstring s) {
    if (s == nullptr) return {};
    const char * c = env->GetStringUTFChars(s, nullptr);
    std::string out = c != nullptr ? c : "";
    if (c != nullptr) env->ReleaseStringUTFChars(s, c);
    return out;
}

ggml_type parseKvType(const std::string & s, ggml_type fallback) {
    if (s == "f16") return GGML_TYPE_F16;
    if (s == "f32") return GGML_TYPE_F32;
    if (s == "q8_0") return GGML_TYPE_Q8_0;
    if (s == "q4_0") return GGML_TYPE_Q4_0;
    if (s == "q4_1") return GGML_TYPE_Q4_1;
    return fallback;
}

// Abort callback for llama_set_abort_callback.
bool abortCb(void * user) {
    auto * e = static_cast<Engine *>(user);
    return e->abort.load(std::memory_order_relaxed);
}

}  // namespace

extern "C" {

JNIEXPORT jlong JNICALL
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeNewFromFile(
        JNIEnv * env, jclass /*clazz*/,
        jstring jPath, jint ctxSize, jint nThreads,
        jstring jTypeK, jstring jTypeV) {

    const std::string path = jstringToUtf8(env, jPath);
    if (path.empty()) {
        throwJava(env, "nativeNewFromFile: empty model path");
        return 0;
    }
    if (access(path.c_str(), R_OK) != 0) {
        throwJava(env, ("nativeNewFromFile: model file not readable: " + path).c_str());
        return 0;
    }

    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers = 0;      // CPU only in milestone 1
    mparams.load_mode = LLAMA_LOAD_MODE_MMAP;  // page cache can be reclaimed under memory pressure

    llama_model * model = llama_model_load_from_file(path.c_str(), mparams);
    if (model == nullptr) {
        throwJava(env, ("nativeNewFromFile: llama failed to load GGUF: " + path).c_str());
        return 0;
    }

    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx = ctxSize > 0 ? (uint32_t) ctxSize : 8192;
    cparams.n_threads = nThreads > 0 ? nThreads : 4;
    cparams.n_threads_batch = nThreads > 0 ? nThreads : 4;
    cparams.type_k = parseKvType(jstringToUtf8(env, jTypeK), GGML_TYPE_Q8_0);
    cparams.type_v = parseKvType(jstringToUtf8(env, jTypeV), GGML_TYPE_F16);

    llama_context * ctx = llama_init_from_model(model, cparams);
    if (ctx == nullptr) {
        llama_model_free(model);
        throwJava(env, "nativeNewFromFile: failed to create llama_context");
        return 0;
    }

    auto * e = new Engine();
    e->model = model;
    e->ctx = ctx;
    e->vocab = llama_model_get_vocab(model);
    llama_set_abort_callback(ctx, abortCb, e);

    LOGI("model loaded: %s n_ctx=%d type_k=%d type_v=%d", path.c_str(),
         cparams.n_ctx, (int) cparams.type_k, (int) cparams.type_v);
    return (jlong) (uintptr_t) e;
}

JNIEXPORT jstring JNICALL
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeGetStringMeta(
        JNIEnv * env, jclass /*clazz*/, jlong ptr, jstring jKey) {

    auto * e = (Engine *) (uintptr_t) ptr;
    if (e == nullptr || e->model == nullptr) return nullptr;
    const std::string key = jstringToUtf8(env, jKey);
    char buf[8192];
    int32_t n = llama_model_meta_val_str(e->model, key.c_str(), buf, sizeof(buf));
    if (n < 0) return nullptr;
    return env->NewStringUTF(buf);
}

// M2:录音中降线程(给 ASR 让核),空闲/成稿时拉满。与在途 generate 共用
// 一把 mu:generate 正在跑时本调用会排队到它结束,不会半路改线程数。
JNIEXPORT void JNICALL
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeSetThreads(
        JNIEnv * /*env*/, jclass /*clazz*/, jlong ptr, jint n) {
    auto * e = (Engine *) (uintptr_t) ptr;
    if (e == nullptr || e->ctx == nullptr || n <= 0) return;
    std::lock_guard<std::mutex> lock(e->mu);
    llama_set_n_threads(e->ctx, n, n);
}

// Blocking generation; call from a background thread. Returns generated text
// (without the prompt). Stops on EOS, maxNewTokens or nativeAbort().
JNIEXPORT jstring JNICALL
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeGenerate(
        JNIEnv * env, jclass /*clazz*/, jlong ptr,
        jstring jPrompt, jint maxNewTokens, jfloat temperature, jint topK, jfloat topP,
        jfloat repeatPenalty, jint seed, jint minNew) {

    auto * e = (Engine *) (uintptr_t) ptr;
    if (e == nullptr || e->ctx == nullptr) {
        throwJava(env, "nativeGenerate: engine is released");
        return nullptr;
    }
    std::lock_guard<std::mutex> lock(e->mu);
    e->abort.store(false);

    // KV 清除(2026-09-16 根治):原实现不清 KV,同引擎连续 generate 会复读
    // 上一条输出 -> Kotlin 侧被迫"每次尝试重建引擎"(E2B 每次重载 4~6s,
    // 是翻译单句 30~40s 的主因)。每次生成前清空内存,引擎即可安全复用;
    // 顺带修复纪要流水线多块共享引擎的潜在污染。
    llama_memory_clear(llama_get_memory(e->ctx), true);

    const std::string prompt = jstringToUtf8(env, jPrompt);
    if (prompt.empty()) {
        throwJava(env, "nativeGenerate: empty prompt");
        return nullptr;
    }

    // --- tokenize the prompt ---
    // parse_special=true: the Kotlin side renders the Gemma 4 template with
    // literal <|turn>/<turn|> markers, which must become control tokens 105/106.
    std::vector<llama_token> tokens(prompt.size() / 2 + 64);
    int32_t n;
    while ((n = llama_tokenize(e->vocab, prompt.data(), (int32_t) prompt.size(),
                               tokens.data(), (int32_t) tokens.size(),
                               /*add_special=*/true, /*parse_special=*/true)) < 0) {
        tokens.resize((size_t) (-n) + 32);
        if (tokens.size() > 4 * 1024 * 1024) {
            throwJava(env, "nativeGenerate: prompt too large to tokenize");
            return nullptr;
        }
    }

    // --- fit into the context window (left-truncate keeps the newest content,
    //     which for our rolling-notes design is the part that must survive) ---
    const uint32_t nCtx = llama_n_ctx(e->ctx);
    const int32_t keep = (int32_t) (nCtx > 64 ? nCtx - 32 : nCtx);
    if (n > keep) {
        LOGE("prompt tokens %d exceed ctx %u, left-truncating", n, nCtx);
        int32_t drop = n - keep;
        // never drop into the middle of multi-byte UTF-8: only relevant on the
        // decode side; tokens are atomic here, so a plain shift is safe.
        tokens.erase(tokens.begin(), tokens.begin() + drop);
        n = keep;
    }

    // --- prefill in chunks ---
    const int32_t chunk = 512;
    llama_pos pos = 0;
    for (int32_t i = 0; i < n; i += chunk) {
        if (e->abort.load()) return env->NewStringUTF("");
        int32_t len = std::min(chunk, n - i);
        llama_batch b = llama_batch_get_one(tokens.data() + i, len);
        if (llama_decode(e->ctx, b) != 0) {
            throwJava(env, "nativeGenerate: prefill decode failed");
            return nullptr;
        }
        pos += len;
    }

    // --- decode loop ---
    const llama_token eos = llama_vocab_eos(e->vocab);
    // seed:调用方显式给随机种子(翻译路径每次点击换盐、每次重试错位,让重试
    // 是真正的再抽签);seed<=0 退回旧行为(42+温度错位,摘要路径用)
    std::mt19937 rng(seed > 0 ? (uint32_t) seed
                              : (42 + (uint32_t) (temperature * 1000.0f)));
    std::string out;
    out.reserve(1024);
    std::vector<char> piece(64);
    int32_t cur = tokens.back();
    const int maxNew = maxNewTokens > 0 ? maxNewTokens : 512;
    const bool greedy = !(temperature > 0.0f);
    const int k = topK > 0 ? topK : 64;          // Gemma 4 recommended: 64
    const float p = topP > 0.0f ? topP : 0.95f;  // recommended: ~0.95
    // Repeat penalty (llama.cpp style, unique-token window): 1B 模型在垃圾/
    // 歧义输入上极易进入复读循环。pen<=0 时禁用(摘要路径保持原行为),
    // 翻译路径传 1.1。只惩罚本步已生成过的 token,窗口=最近 64 个生成 token。
    const bool usePen = repeatPenalty > 1.0f;
    std::vector<llama_token> recent;
    std::vector<bool> penalized;  // scratch for unique marking

    for (int i = 0; i < maxNew; ++i) {
        if (e->abort.load()) break;
        llama_batch b = llama_batch_get_one(&cur, 1);
        if (llama_decode(e->ctx, b) != 0) {
            throwJava(env, "nativeGenerate: decode failed");
            return nullptr;
        }
        float * logits = llama_get_logits_ith(e->ctx, 0);
        const int vsize = llama_vocab_n_tokens(e->vocab);

        // 防空白输出(2026-09-16 真机实锤:极端量化下个别输入首 token 直落
        // EOS,表现为 raw.len=0 的"空输出"):前 minNew 步把 EOS 置 -inf,
        // 强制模型至少吐几个 token;有害输出交给上层质检闸门。
        // minNew<=0 = 关闭(摘要路径用)。
        if (minNew > 0 && i < minNew && eos >= 0 && eos < vsize) {
            logits[eos] = -INFINITY;
        }

        if (usePen && !recent.empty()) {
            if ((int) recent.size() > 64)
                recent.erase(recent.begin(), recent.begin() + (recent.size() - 64));
            if (penalized.size() < (size_t) vsize)
                penalized.assign((size_t) vsize + 4096, false);
            for (llama_token t : recent) {
                if (t >= 0 && t < vsize && !penalized[t]) {
                    penalized[t] = true;
                    float & lg = logits[t];
                    lg = lg < 0 ? lg * repeatPenalty : lg / repeatPenalty;
                }
            }
            for (llama_token t : recent) if (t >= 0 && t < vsize) penalized[t] = false;
        }

        // pick next token
        llama_token next;
        if (greedy) {
            next = 0;
            float best = logits[0];
            for (int t = 1; t < vsize; ++t) {
                if (logits[t] > best) { best = logits[t]; next = t; }
            }
        } else {
            // top-k cut (partial sort), then temperature softmax + top-p cut, then sample
            std::vector<std::pair<float, int32_t>> cand;
            cand.reserve(k);
            for (int t = 0; t < vsize; ++t) {
                if (cand.size() < (size_t) k) {
                    cand.emplace_back(logits[t], t);
                    std::push_heap(cand.begin(), cand.end());
                } else if (logits[t] > cand.front().first) {
                    std::pop_heap(cand.begin(), cand.end());
                    cand.back() = {logits[t], t};
                    std::push_heap(cand.begin(), cand.end());
                }
            }
            float maxv = cand.front().first;
            double sum = 0;
            for (auto & c : cand) { c.first = std::exp((c.first - maxv) / temperature); sum += c.first; }
            std::sort(cand.begin(), cand.end(), std::greater<>());
            std::uniform_real_distribution<double> dist(0.0, sum);
            double target = dist(rng), acc = 0;
            next = cand[0].second;
            for (auto & c : cand) {
                acc += c.first;
                next = c.second;
                if (acc / sum >= p) break;
            }
        }

        if (eos >= 0 && next == eos) break;
        if (llama_vocab_is_eog(e->vocab, next)) break;

        int32_t len = llama_token_to_piece(e->vocab, next, piece.data(),
                                           (int32_t) piece.size(), 0, false);
        if (len > 0) {
            out.append(piece.data(), (size_t) len);
            // Gemma 4 closes its turn with "<turn|>" (id 106). It may not be
            // registered as eog in GGUF metadata, so stop on the string too
            // and strip it from the output.
            if (out.size() >= 7 && out.compare(out.size() - 7, 7, "<turn|>") == 0) {
                out.erase(out.size() - 7);
                break;
            }
        }
        cur = next;
        if (usePen) recent.push_back(next);
    }

    LOGI("generated %d chars", (int) out.size());
    return env->NewStringUTF(out.c_str());
}

JNIEXPORT void JNICALL
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeAbort(JNIEnv * /*env*/, jclass, jlong ptr) {
    auto * e = (Engine *) (uintptr_t) ptr;
    if (e != nullptr) e->abort.store(true);
}

JNIEXPORT void JNICALL
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeFree(JNIEnv * /*env*/, jclass, jlong ptr) {
    auto * e = (Engine *) (uintptr_t) ptr;
    if (e == nullptr) return;
    std::lock_guard<std::mutex> lock(e->mu);  // wait for an in-flight generate
    if (e->ctx != nullptr) llama_free(e->ctx);
    if (e->model != nullptr) llama_model_free(e->model);
    delete e;
    LOGI("engine released");
}

}  // extern "C"