// Persistent RKLLM inference server for the CallRec app.
// Loads a .rkllm model onto the RK3588 NPU ONCE, then answers prompts over a
// simple line protocol on stdin/stdout so Node can spawn it a single time:
//   stdin : one request per line = base64(prompt)          (base64 has no newlines)
//   stdout: "READY" once after model load, then per request:
//           "OK <base64(completion)>"  on success
//           "ERR <message>"            on failure
// keep_history=0 → every request is independent (Node sends full context each time).
#include <cstddef>   // rkllm.h uses size_t without including it itself
#include "rkllm.h"
#include <cstdio>
#include <cstring>
#include <string>
#include <iostream>

static std::string g_out;
static bool g_error = false;

static int cb(RKLLMResult* r, void* /*ud*/, LLMCallState st) {
    if (st == RKLLM_RUN_NORMAL) { if (r && r->text) g_out += r->text; }
    else if (st == RKLLM_RUN_ERROR) { g_error = true; }
    return 0;
}

static const char B64[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
static std::string b64enc(const std::string& in) {
    std::string out; int val = 0, bits = -6;
    for (unsigned char c : in) {
        val = (val << 8) + c; bits += 8;
        while (bits >= 0) { out.push_back(B64[(val >> bits) & 0x3F]); bits -= 6; }
    }
    if (bits > -6) out.push_back(B64[((val << 8) >> (bits + 8)) & 0x3F]);
    while (out.size() % 4) out.push_back('=');
    return out;
}
static std::string b64dec(const std::string& in) {
    int T[256]; for (int i = 0; i < 256; i++) T[i] = -1;
    for (int i = 0; i < 64; i++) T[(unsigned char)B64[i]] = i;
    std::string out; int val = 0, bits = -8;
    for (unsigned char c : in) {
        if (T[c] == -1) continue;
        val = (val << 6) + T[c]; bits += 6;
        if (bits >= 0) { out.push_back((char)((val >> bits) & 0xFF)); bits -= 8; }
    }
    return out;
}

int main(int argc, char** argv) {
    if (argc < 2) { fprintf(stderr, "usage: %s model.rkllm\n", argv[0]); return 1; }
    RKLLMParam p = rkllm_createDefaultParam();
    p.model_path        = argv[1];
    p.max_context_len   = 8192;
    p.max_new_tokens    = 1536;
    // Gemma-3-4B (w8a8) degenerates into repetition loops ("X, X, X…") and echoes
    // the prompt under pure-greedy decode → needs real sampling + strong repeat/
    // frequency penalties. (Qwen was fine at top_k=1/1.1; Gemma is not.)
    p.top_k             = 20;
    p.top_p             = 0.9f;
    p.temperature       = 0.5f;
    p.repeat_penalty    = 1.3f;
    p.frequency_penalty = 0.5f;
    p.is_async          = false;  // rkllm_run blocks until finish → buffer complete on return
    // Pin the runtime's CPU-side threads to the big A76 cores (4-7). The 1.2.1
    // default mask targets the little A55 cores, which starves de/quant and
    // sampling between NPU calls.
    p.extend_param.enabled_cpus_num  = 4;
    p.extend_param.enabled_cpus_mask = (1u<<4)|(1u<<5)|(1u<<6)|(1u<<7);
    p.skip_special_token = true;

    LLMHandle h = nullptr;
    RKLLMCallback cbs; memset(&cbs, 0, sizeof(cbs)); cbs.result_callback = cb;
    if (rkllm_init(&h, &p, &cbs) != 0) { fprintf(stderr, "rkllm_init failed\n"); return 2; }
    // Gemma-3 turn format. rkllm does NOT auto-apply the HF chat template, so
    // without this the model gets the raw prompt with no turn markers and just
    // rambles generic filler. prefix/postfix wrap each user prompt; <bos> is
    // added by the runtime (add_bos_token=True at conversion).
    rkllm_set_chat_template(h, "", "<start_of_turn>user\n", "<end_of_turn>\n<start_of_turn>model\n");
    printf("READY\n"); fflush(stdout);
    fprintf(stderr, "rkllm model loaded\n"); fflush(stderr);

    std::string line;
    while (std::getline(std::cin, line)) {
        if (line.empty()) continue;
        std::string prompt = b64dec(line);
        g_out.clear(); g_error = false;

        RKLLMInput in; memset(&in, 0, sizeof(in));
        in.role         = "user";
        in.enable_thinking = false;
        in.input_type   = RKLLM_INPUT_PROMPT;
        in.prompt_input = prompt.c_str();

        RKLLMInferParam ip; memset(&ip, 0, sizeof(ip));
        ip.mode         = RKLLM_INFER_GENERATE;
        ip.keep_history = 0;

        rkllm_clear_kv_cache(h, 0, nullptr, nullptr);  // hard reset: no cross-recording KV bleed
        int rc = rkllm_run(h, &in, &ip, nullptr);
        if (rc != 0 || g_error) { printf("ERR run rc=%d\n", rc); fflush(stdout); continue; }
        printf("OK %s\n", b64enc(g_out).c_str()); fflush(stdout);
    }
    rkllm_destroy(h);
    return 0;
}