#include "jllama.h"

#include "arg.h"
#include "json-schema-to-grammar.h"
#include "llama.h"
#include "log.h"
#include "nlohmann/json.hpp"
#include "server.hpp"

#include <functional>
#include <iostream>
#include <stdexcept>

// We store some references to Java classes and their fields/methods here to speed up things for later and to fail
// early on if anything can't be found. This happens when the JVM loads the shared library (see `JNI_OnLoad`).
// The references remain valid throughout the whole life of the shared library, on `JNI_OnUnload` they are released.

namespace {
JavaVM *g_vm = nullptr;

// classes
jclass c_llama_model = nullptr;
jclass c_llama_iterator = nullptr;
jclass c_standard_charsets = nullptr;
jclass c_output = nullptr;
jclass c_string = nullptr;
jclass c_hash_map = nullptr;
jclass c_map = nullptr;
jclass c_set = nullptr;
jclass c_entry = nullptr;
jclass c_iterator = nullptr;
jclass c_integer = nullptr;
jclass c_float = nullptr;
jclass c_biconsumer = nullptr;
jclass c_llama_error = nullptr;
jclass c_log_level = nullptr;
jclass c_log_format = nullptr;
jclass c_error_oom = nullptr;

// constructors
jmethodID cc_output = nullptr;
jmethodID cc_hash_map = nullptr;
jmethodID cc_integer = nullptr;
jmethodID cc_float = nullptr;

// methods
jmethodID m_get_bytes = nullptr;
jmethodID m_entry_set = nullptr;
jmethodID m_set_iterator = nullptr;
jmethodID m_iterator_has_next = nullptr;
jmethodID m_iterator_next = nullptr;
jmethodID m_entry_key = nullptr;
jmethodID m_entry_value = nullptr;
jmethodID m_map_put = nullptr;
jmethodID m_int_value = nullptr;
jmethodID m_float_value = nullptr;
jmethodID m_biconsumer_accept = nullptr;

// fields
jfieldID f_model_pointer = nullptr;
jfieldID f_task_id = nullptr;
jfieldID f_utf_8 = nullptr;
jfieldID f_iter_has_next = nullptr;
jfieldID f_log_level_debug = nullptr;
jfieldID f_log_level_info = nullptr;
jfieldID f_log_level_warn = nullptr;
jfieldID f_log_level_error = nullptr;
jfieldID f_log_format_json = nullptr;
jfieldID f_log_format_text = nullptr;

// objects
jobject o_utf_8 = nullptr;
jobject o_log_level_debug = nullptr;
jobject o_log_level_info = nullptr;
jobject o_log_level_warn = nullptr;
jobject o_log_level_error = nullptr;
jobject o_log_format_json = nullptr;
jobject o_log_format_text = nullptr;
jobject o_log_callback = nullptr;

/**
 * Convert a Java string to a std::string
 */
std::string parse_jstring(JNIEnv *env, jstring java_string) {
    auto *const string_bytes = (jbyteArray)env->CallObjectMethod(java_string, m_get_bytes, o_utf_8);

    auto length = (size_t)env->GetArrayLength(string_bytes);
    jbyte *byte_elements = env->GetByteArrayElements(string_bytes, nullptr);

    std::string string = std::string((char *)byte_elements, length);

    env->ReleaseByteArrayElements(string_bytes, byte_elements, JNI_ABORT);
    env->DeleteLocalRef(string_bytes);

    return string;
}

char **parse_string_array(JNIEnv *env, const jobjectArray string_array, const jsize length) {
    auto *const result = static_cast<char **>(malloc(length * sizeof(char *)));

    if (result == nullptr) {
        return nullptr;
    }

    for (jsize i = 0; i < length; i++) {
        auto *const javaString = static_cast<jstring>(env->GetObjectArrayElement(string_array, i));
        const char *cString = env->GetStringUTFChars(javaString, nullptr);
        result[i] = strdup(cString);
        env->ReleaseStringUTFChars(javaString, cString);
    }

    return result;
}

void free_string_array(char **array, jsize length) {
    if (array != nullptr) {
        for (jsize i = 0; i < length; i++) {
            free(array[i]);
        }
        free(array);
    }
}

/**
 * Since Java expects utf16 but std::strings are utf8, we can't directly use `env->NewString` or `env-NewString`,
 * but we directly send the bytes and do the conversion in Java. Unfortunately, there isn't a nice/standardized way to
 * do this conversion in C++
 */
jbyteArray parse_jbytes(JNIEnv *env, const std::string &string) {
    jsize length = string.size(); // NOLINT(*-narrowing-conversions)
    jbyteArray bytes = env->NewByteArray(length);
    env->SetByteArrayRegion(bytes, 0, length, reinterpret_cast<const jbyte *>(string.c_str()));
    return bytes;
}

/**
 * Map a llama.cpp log level to its Java enumeration option.
 */
jobject log_level_to_jobject(ggml_log_level level) {
    switch (level) {
    case GGML_LOG_LEVEL_ERROR:
        return o_log_level_error;
    case GGML_LOG_LEVEL_WARN:
        return o_log_level_warn;
    default:
    case GGML_LOG_LEVEL_INFO:
        return o_log_level_info;
    case GGML_LOG_LEVEL_DEBUG:
        return o_log_level_debug;
    }
}

/**
 * Returns the JNIEnv of the current thread.
 */
JNIEnv *get_jni_env() {
    JNIEnv *env = nullptr;
    if (g_vm == nullptr || g_vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6) != JNI_OK) {
        throw std::runtime_error("Thread is not attached to the JVM");
    }
    return env;
}

bool log_json;
std::function<void(ggml_log_level, const char *, void *)> log_callback;

/**
 * Invoke the log callback if there is any.
 */
void log_callback_trampoline(ggml_log_level level, const char *text, void *user_data) {
    if (log_callback != nullptr) {
        log_callback(level, text, user_data);
    }
}
} // namespace

/**
 * The VM calls JNI_OnLoad when the native library is loaded (for example, through `System.loadLibrary`).
 * `JNI_OnLoad` must return the JNI version needed by the native library.
 * In order to use any of the new JNI functions, a native library must export a `JNI_OnLoad` function that returns
 * `JNI_VERSION_1_2`. If the native library does not export a JNI_OnLoad function, the VM assumes that the library
 * only requires JNI version `JNI_VERSION_1_1`. If the VM does not recognize the version number returned by
 `JNI_OnLoad`, the VM will unload the library and act as if the library was never loaded.
 */
JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *reserved) {
    g_vm = vm;
    JNIEnv *env = nullptr;

    if (JNI_OK != vm->GetEnv((void **)&env, JNI_VERSION_1_1)) {
        goto error;
    }

    // find classes
    c_llama_model = env->FindClass("de/kherud/llama/LlamaModel");
    c_llama_iterator = env->FindClass("de/kherud/llama/LlamaIterator");
    c_standard_charsets = env->FindClass("java/nio/charset/StandardCharsets");
    c_output = env->FindClass("de/kherud/llama/LlamaOutput");
    c_string = env->FindClass("java/lang/String");
    c_hash_map = env->FindClass("java/util/HashMap");
    c_map = env->FindClass("java/util/Map");
    c_set = env->FindClass("java/util/Set");
    c_entry = env->FindClass("java/util/Map$Entry");
    c_iterator = env->FindClass("java/util/Iterator");
    c_integer = env->FindClass("java/lang/Integer");
    c_float = env->FindClass("java/lang/Float");
    c_biconsumer = env->FindClass("java/util/function/BiConsumer");
    c_llama_error = env->FindClass("de/kherud/llama/LlamaException");
    c_log_level = env->FindClass("de/kherud/llama/LogLevel");
    c_log_format = env->FindClass("de/kherud/llama/args/LogFormat");
    c_error_oom = env->FindClass("java/lang/OutOfMemoryError");

    if (!(c_llama_model && c_llama_iterator && c_standard_charsets && c_output && c_string && c_hash_map && c_map &&
          c_set && c_entry && c_iterator && c_integer && c_float && c_biconsumer && c_llama_error && c_log_level &&
          c_log_format && c_error_oom)) {
        goto error;
    }

    // create references
    c_llama_model = (jclass)env->NewGlobalRef(c_llama_model);
    c_llama_iterator = (jclass)env->NewGlobalRef(c_llama_iterator);
    c_output = (jclass)env->NewGlobalRef(c_output);
    c_string = (jclass)env->NewGlobalRef(c_string);
    c_hash_map = (jclass)env->NewGlobalRef(c_hash_map);
    c_map = (jclass)env->NewGlobalRef(c_map);
    c_set = (jclass)env->NewGlobalRef(c_set);
    c_entry = (jclass)env->NewGlobalRef(c_entry);
    c_iterator = (jclass)env->NewGlobalRef(c_iterator);
    c_integer = (jclass)env->NewGlobalRef(c_integer);
    c_float = (jclass)env->NewGlobalRef(c_float);
    c_biconsumer = (jclass)env->NewGlobalRef(c_biconsumer);
    c_llama_error = (jclass)env->NewGlobalRef(c_llama_error);
    c_log_level = (jclass)env->NewGlobalRef(c_log_level);
    c_log_format = (jclass)env->NewGlobalRef(c_log_format);
    c_error_oom = (jclass)env->NewGlobalRef(c_error_oom);

    // find constructors
    cc_output = env->GetMethodID(c_output, "<init>", "([BLjava/util/Map;Z)V");
    cc_hash_map = env->GetMethodID(c_hash_map, "<init>", "()V");
    cc_integer = env->GetMethodID(c_integer, "<init>", "(I)V");
    cc_float = env->GetMethodID(c_float, "<init>", "(F)V");

    if (!(cc_output && cc_hash_map && cc_integer && cc_float)) {
        goto error;
    }

    // find methods
    m_get_bytes = env->GetMethodID(c_string, "getBytes", "(Ljava/lang/String;)[B");
    m_entry_set = env->GetMethodID(c_map, "entrySet", "()Ljava/util/Set;");
    m_set_iterator = env->GetMethodID(c_set, "iterator", "()Ljava/util/Iterator;");
    m_iterator_has_next = env->GetMethodID(c_iterator, "hasNext", "()Z");
    m_iterator_next = env->GetMethodID(c_iterator, "next", "()Ljava/lang/Object;");
    m_entry_key = env->GetMethodID(c_entry, "getKey", "()Ljava/lang/Object;");
    m_entry_value = env->GetMethodID(c_entry, "getValue", "()Ljava/lang/Object;");
    m_map_put = env->GetMethodID(c_map, "put", "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");
    m_int_value = env->GetMethodID(c_integer, "intValue", "()I");
    m_float_value = env->GetMethodID(c_float, "floatValue", "()F");
    m_biconsumer_accept = env->GetMethodID(c_biconsumer, "accept", "(Ljava/lang/Object;Ljava/lang/Object;)V");

    if (!(m_get_bytes && m_entry_set && m_set_iterator && m_iterator_has_next && m_iterator_next && m_entry_key &&
          m_entry_value && m_map_put && m_int_value && m_float_value && m_biconsumer_accept)) {
        goto error;
    }

    // find fields
    f_model_pointer = env->GetFieldID(c_llama_model, "ctx", "J");
    f_task_id = env->GetFieldID(c_llama_iterator, "taskId", "I");
    f_utf_8 = env->GetStaticFieldID(c_standard_charsets, "UTF_8", "Ljava/nio/charset/Charset;");
    f_iter_has_next = env->GetFieldID(c_llama_iterator, "hasNext", "Z");
    f_log_level_debug = env->GetStaticFieldID(c_log_level, "DEBUG", "Lde/kherud/llama/LogLevel;");
    f_log_level_info = env->GetStaticFieldID(c_log_level, "INFO", "Lde/kherud/llama/LogLevel;");
    f_log_level_warn = env->GetStaticFieldID(c_log_level, "WARN", "Lde/kherud/llama/LogLevel;");
    f_log_level_error = env->GetStaticFieldID(c_log_level, "ERROR", "Lde/kherud/llama/LogLevel;");
    f_log_format_json = env->GetStaticFieldID(c_log_format, "JSON", "Lde/kherud/llama/args/LogFormat;");
    f_log_format_text = env->GetStaticFieldID(c_log_format, "TEXT", "Lde/kherud/llama/args/LogFormat;");

    if (!(f_model_pointer && f_task_id && f_utf_8 && f_iter_has_next && f_log_level_debug && f_log_level_info &&
          f_log_level_warn && f_log_level_error && f_log_format_json && f_log_format_text)) {
        goto error;
    }

    o_utf_8 = env->NewStringUTF("UTF-8");
    o_log_level_debug = env->GetStaticObjectField(c_log_level, f_log_level_debug);
    o_log_level_info = env->GetStaticObjectField(c_log_level, f_log_level_info);
    o_log_level_warn = env->GetStaticObjectField(c_log_level, f_log_level_warn);
    o_log_level_error = env->GetStaticObjectField(c_log_level, f_log_level_error);
    o_log_format_json = env->GetStaticObjectField(c_log_format, f_log_format_json);
    o_log_format_text = env->GetStaticObjectField(c_log_format, f_log_format_text);

    if (!(o_utf_8 && o_log_level_debug && o_log_level_info && o_log_level_warn && o_log_level_error &&
          o_log_format_json && o_log_format_text)) {
        goto error;
    }

    o_utf_8 = env->NewGlobalRef(o_utf_8);
    o_log_level_debug = env->NewGlobalRef(o_log_level_debug);
    o_log_level_info = env->NewGlobalRef(o_log_level_info);
    o_log_level_warn = env->NewGlobalRef(o_log_level_warn);
    o_log_level_error = env->NewGlobalRef(o_log_level_error);
    o_log_format_json = env->NewGlobalRef(o_log_format_json);
    o_log_format_text = env->NewGlobalRef(o_log_format_text);

    if (env->ExceptionCheck()) {
        env->ExceptionDescribe();
        goto error;
    }

    llama_backend_init();

    goto success;

error:
    return JNI_ERR;

success:
    return JNI_VERSION_1_6;
}

/**
 * The VM calls `JNI_OnUnload` when the class loader containing the native library is garbage collected.
 * This function can be used to perform cleanup operations. Because this function is called in an unknown context
 * (such as from a finalizer), the programmer should be conservative on using Java VM services, and refrain from
 * arbitrary Java call-backs.
 * Note that `JNI_OnLoad` and `JNI_OnUnload` are two functions optionally supplied by JNI libraries, not exported from
 * the VM.
 */
JNIEXPORT void JNICALL JNI_OnUnload(JavaVM *vm, void *reserved) {
    JNIEnv *env = nullptr;

    if (JNI_OK != vm->GetEnv((void **)&env, JNI_VERSION_1_6)) {
        return;
    }

    env->DeleteGlobalRef(c_llama_model);
    env->DeleteGlobalRef(c_llama_iterator);
    env->DeleteGlobalRef(c_output);
    env->DeleteGlobalRef(c_string);
    env->DeleteGlobalRef(c_hash_map);
    env->DeleteGlobalRef(c_map);
    env->DeleteGlobalRef(c_set);
    env->DeleteGlobalRef(c_entry);
    env->DeleteGlobalRef(c_iterator);
    env->DeleteGlobalRef(c_integer);
    env->DeleteGlobalRef(c_float);
    env->DeleteGlobalRef(c_biconsumer);
    env->DeleteGlobalRef(c_llama_error);
    env->DeleteGlobalRef(c_log_level);
    env->DeleteGlobalRef(c_log_level);
    env->DeleteGlobalRef(c_error_oom);

    env->DeleteGlobalRef(o_utf_8);
    env->DeleteGlobalRef(o_log_level_debug);
    env->DeleteGlobalRef(o_log_level_info);
    env->DeleteGlobalRef(o_log_level_warn);
    env->DeleteGlobalRef(o_log_level_error);
    env->DeleteGlobalRef(o_log_format_json);
    env->DeleteGlobalRef(o_log_format_text);

    if (o_log_callback != nullptr) {
        env->DeleteGlobalRef(o_log_callback);
    }

    llama_backend_free();
}

JNIEXPORT void JNICALL Java_de_kherud_llama_LlamaModel_loadModel(JNIEnv *env, jobject obj, jobjectArray jparams) {
    common_params params;

    const jsize argc = env->GetArrayLength(jparams);
    char **argv = parse_string_array(env, jparams, argc);
    if (argv == nullptr) {
        return;
    }

    const auto parsed_params = common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER);
    free_string_array(argv, argc);
    if (!parsed_params) {
        return;
    }

    SRV_INF("loading model '%s'\n", params.model.c_str());

    common_init();

    // struct that contains llama context and inference
    auto *ctx_server = new server_context();

    llama_numa_init(params.numa);

    LOG_INF("system info: n_threads = %d, n_threads_batch = %d, total_threads = %d\n", params.cpuparams.n_threads,
            params.cpuparams_batch.n_threads, std::thread::hardware_concurrency());
    LOG_INF("\n");
    LOG_INF("%s\n", common_params_get_system_info(params).c_str());
    LOG_INF("\n");

    std::atomic<server_state> state{SERVER_STATE_LOADING_MODEL};

    // Necessary similarity of prompt for slot selection
    ctx_server->slot_prompt_similarity = params.slot_prompt_similarity;

    LOG_INF("%s: loading model\n", __func__);

    // load the model
    if (!ctx_server->load_model(params)) {
        llama_backend_free();
        env->ThrowNew(c_llama_error, "could not load model from given file path");
        return;
    }

    ctx_server->init();
    state.store(SERVER_STATE_READY);

    LOG_INF("%s: model loaded\n", __func__);

    const auto model_meta = ctx_server->model_meta();

    if (!params.speculative.model.empty() || !params.speculative.hf_repo.empty()) {
        SRV_INF("loading draft model '%s'\n", params.speculative.model.c_str());
        auto params_dft = params;

        params_dft.devices = params.speculative.devices;
        params_dft.hf_file = params.speculative.hf_file;
        params_dft.hf_repo = params.speculative.hf_repo;
        params_dft.model = params.speculative.model;
        params_dft.model_url = params.speculative.model_url;
        params_dft.n_ctx = params.speculative.n_ctx == 0 ? params.n_ctx / params.n_parallel : params.speculative.n_ctx;
        params_dft.n_gpu_layers = params.speculative.n_gpu_layers;
        params_dft.n_parallel = 1;

        common_init_result llama_init_dft = common_init_from_params(params_dft);

        llama_model *model_dft = llama_init_dft.model.get();

        if (model_dft == nullptr) {
            SRV_ERR("failed to load draft model, '%s'\n", params.speculative.model.c_str());
        }

        if (!common_speculative_are_compatible(ctx_server->ctx, llama_init_dft.context.get())) {
            SRV_ERR("the draft model '%s' is not compatible with the target model '%s'\n",
                    params.speculative.model.c_str(), params.model.c_str());
        }

        const int n_ctx_dft = llama_n_ctx(llama_init_dft.context.get());

        ctx_server->cparams_dft = common_context_params_to_llama(params_dft);
        ctx_server->cparams_dft.n_batch = n_ctx_dft;

        // force F16 KV cache for the draft model for extra performance
        ctx_server->cparams_dft.type_k = GGML_TYPE_F16;
        ctx_server->cparams_dft.type_v = GGML_TYPE_F16;

        // the context is not needed - we will create one for each slot
        llama_init_dft.context.reset();
    }

    // print sample chat example to make it clear which template is used
    LOG_INF("%s: chat template, chat_template: %s, example_format: '%s'\n", __func__,
            common_chat_templates_source(ctx_server->chat_templates.get()),
            common_chat_format_example(ctx_server->chat_templates.get(), ctx_server->params_base.use_jinja).c_str());

    // print sample chat example to make it clear which template is used
    //    LOG_INF("%s: chat template, chat_template: %s, example_format: '%s'\n", __func__,
    //         common_chat_templates_source(ctx_server->chat_templates.get()),
    //        common_chat_format_example(*ctx_server->chat_templates.template_default,
    //        ctx_server->params_base.use_jinja) .c_str());

    ctx_server->queue_tasks.on_new_task(
        std::bind(&server_context::process_single_task, ctx_server, std::placeholders::_1));
    ctx_server->queue_tasks.on_update_slots(std::bind(&server_context::update_slots, ctx_server));

    std::thread t([ctx_server]() {
        JNIEnv *env;
        jint res = g_vm->GetEnv((void **)&env, JNI_VERSION_1_6);
        if (res == JNI_EDETACHED) {
            res = g_vm->AttachCurrentThread((void **)&env, nullptr);
            if (res != JNI_OK) {
                throw std::runtime_error("Failed to attach thread to JVM");
            }
        }
        ctx_server->queue_tasks.start_loop();
    });
    t.detach();

    env->SetLongField(obj, f_model_pointer, reinterpret_cast<jlong>(ctx_server));
}

JNIEXPORT jint JNICALL Java_de_kherud_llama_LlamaModel_requestCompletion(JNIEnv *env, jobject obj, jstring jparams) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    std::string c_params = parse_jstring(env, jparams);
    json data = json::parse(c_params);

    server_task_type type = SERVER_TASK_TYPE_COMPLETION;

    if (data.contains("input_prefix") || data.contains("input_suffix")) {
        type = SERVER_TASK_TYPE_INFILL;
    }

    auto completion_id = gen_chatcmplid();
    std::vector<server_task> tasks;

    try {
        const auto &prompt = data.at("prompt");

        std::vector<llama_tokens> tokenized_prompts = tokenize_input_prompts(ctx_server->vocab, prompt, true, true);

        tasks.reserve(tokenized_prompts.size());
        for (size_t i = 0; i < tokenized_prompts.size(); i++) {
            server_task task = server_task(type);

            task.id = ctx_server->queue_tasks.get_new_id();
            task.index = i;

            task.prompt_tokens = std::move(tokenized_prompts[i]);
            task.params = server_task::params_from_json_cmpl(ctx_server->ctx, ctx_server->params_base, data);
            task.id_selected_slot = json_value(data, "id_slot", -1);

            // OAI-compat
            task.params.oaicompat = OAICOMPAT_TYPE_NONE;
            task.params.oaicompat_cmpl_id = completion_id;
            // oaicompat_model is already populated by params_from_json_cmpl

            tasks.push_back(task);
        }
    } catch (const std::exception &e) {
        const auto &err = format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST);
        env->ThrowNew(c_llama_error, err.dump().c_str());
        return 0;
    }

    ctx_server->queue_results.add_waiting_tasks(tasks);
    ctx_server->queue_tasks.post(tasks);

    const auto task_ids = server_task::get_list_id(tasks);

    if (task_ids.size() != 1) {
        env->ThrowNew(c_llama_error, "multitasking currently not supported");
        return 0;
    }

    return *task_ids.begin();
}

JNIEXPORT void JNICALL Java_de_kherud_llama_LlamaModel_releaseTask(JNIEnv *env, jobject obj, jint id_task) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)
    ctx_server->queue_results.remove_waiting_task_id(id_task);
}

JNIEXPORT jobject JNICALL Java_de_kherud_llama_LlamaModel_receiveCompletion(JNIEnv *env, jobject obj, jint id_task) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    server_task_result_ptr result = ctx_server->queue_results.recv(id_task);

    if (result->is_error()) {
        std::string response = result->to_json()["message"].get<std::string>();
        ctx_server->queue_results.remove_waiting_task_id(id_task);
        env->ThrowNew(c_llama_error, response.c_str());
        return nullptr;
    }
    const auto out_res = result->to_json();

    std::string response = out_res["content"].get<std::string>();
    if (result->is_stop()) {
        ctx_server->queue_results.remove_waiting_task_id(id_task);
    }

    jobject o_probabilities = env->NewObject(c_hash_map, cc_hash_map);
    if (out_res.contains("completion_probabilities")) {
        auto completion_probabilities = out_res["completion_probabilities"];
        for (const auto &entry : completion_probabilities) {
            auto probs = entry["probs"];
            for (const auto &tp : probs) {
                std::string tok_str = tp["tok_str"];
                jstring jtok_str = env->NewStringUTF(tok_str.c_str());
                float prob = tp["prob"];
                jobject jprob = env->NewObject(c_float, cc_float, prob);
                env->CallObjectMethod(o_probabilities, m_map_put, jtok_str, jprob);
                env->DeleteLocalRef(jtok_str);
                env->DeleteLocalRef(jprob);
            }
        }
    }
    jbyteArray jbytes = parse_jbytes(env, response);
    return env->NewObject(c_output, cc_output, jbytes, o_probabilities, result->is_stop());
}

JNIEXPORT jfloatArray JNICALL Java_de_kherud_llama_LlamaModel_embed(JNIEnv *env, jobject obj, jstring jprompt) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    if (!ctx_server->params_base.embedding) {
        env->ThrowNew(c_llama_error,
                      "model was not loaded with embedding support (see ModelParameters#setEmbedding(boolean))");
        return nullptr;
    }

    const std::string prompt = parse_jstring(env, jprompt);

    SRV_INF("Calling embedding '%s'\n", prompt.c_str());

    const auto tokens = tokenize_mixed(ctx_server->vocab, prompt, true, true);
    std::vector<server_task> tasks;

    server_task task = server_task(SERVER_TASK_TYPE_EMBEDDING);

    task.id = ctx_server->queue_tasks.get_new_id();
    task.index = 0;
    task.prompt_tokens = std::move(tokens);

    // OAI-compat
    task.params.oaicompat = OAICOMPAT_TYPE_NONE;

    tasks.push_back(task);

    ctx_server->queue_results.add_waiting_tasks(tasks);
    ctx_server->queue_tasks.post(tasks);

    std::unordered_set<int> task_ids = server_task::get_list_id(tasks);
    const auto id_task = *task_ids.begin();
    json responses = json::array();

    json error = nullptr;

    server_task_result_ptr result = ctx_server->queue_results.recv(id_task);

    json response_str = result->to_json();
    if (result->is_error()) {
        std::string response = result->to_json()["message"].get<std::string>();
        ctx_server->queue_results.remove_waiting_task_id(id_task);
        env->ThrowNew(c_llama_error, response.c_str());
        return nullptr;
    }

    if (result->is_stop()) {
        ctx_server->queue_results.remove_waiting_task_id(id_task);
    }

    const auto out_res = result->to_json();

    // Extract "embedding" as a vector of vectors (2D array)
    std::vector<std::vector<float>> embedding = out_res["embedding"].get<std::vector<std::vector<float>>>();

    // Get total number of rows in the embedding
    jsize embedding_rows = embedding.size();

    // Get total number of columns in the first row (assuming all rows are of equal length)
    jsize embedding_cols = embedding_rows > 0 ? embedding[0].size() : 0;

    SRV_INF("Embedding has %d rows and %d columns\n", embedding_rows, embedding_cols);

    // Ensure embedding is not empty
    if (embedding.empty() || embedding[0].empty()) {
        env->ThrowNew(c_error_oom, "embedding array is empty");
        return nullptr;
    }

    // Extract only the first row
    const std::vector<float> &first_row = embedding[0]; // Reference to avoid copying

    // Create a new float array in JNI
    jfloatArray j_embedding = env->NewFloatArray(embedding_cols);
    if (j_embedding == nullptr) {
        env->ThrowNew(c_error_oom, "could not allocate embedding");
        return nullptr;
    }

    // Copy the first row into the JNI float array
    env->SetFloatArrayRegion(j_embedding, 0, embedding_cols, reinterpret_cast<const jfloat *>(first_row.data()));

    return j_embedding;
}

JNIEXPORT jobject JNICALL Java_de_kherud_llama_LlamaModel_rerank(JNIEnv *env, jobject obj, jstring jprompt,
                                                                 jobjectArray documents) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    if (!ctx_server->params_base.reranking || ctx_server->params_base.embedding) {
        env->ThrowNew(c_llama_error,
                      "This server does not support reranking. Start it with `--reranking` and without `--embedding`");
        return nullptr;
    }

    const std::string prompt = parse_jstring(env, jprompt);

    const auto tokenized_query = tokenize_mixed(ctx_server->vocab, prompt, true, true);

    json responses = json::array();

    std::vector<server_task> tasks;
    const jsize amount_documents = env->GetArrayLength(documents);
    auto *document_array = parse_string_array(env, documents, amount_documents);
    auto document_vector = std::vector<std::string>(document_array, document_array + amount_documents);
    free_string_array(document_array, amount_documents);

    std::vector<llama_tokens> tokenized_docs = tokenize_input_prompts(ctx_server->vocab, document_vector, true, true);

    tasks.reserve(tokenized_docs.size());
    for (int i = 0; i < tokenized_docs.size(); i++) {
        auto task = server_task(SERVER_TASK_TYPE_RERANK);
        task.id = ctx_server->queue_tasks.get_new_id();
        task.index = i;
        task.prompt_tokens = format_rerank(ctx_server->vocab, tokenized_query, tokenized_docs[i]);
        tasks.push_back(task);
    }
    ctx_server->queue_results.add_waiting_tasks(tasks);
    ctx_server->queue_tasks.post(tasks);

    // get the result
    std::unordered_set<int> task_ids = server_task::get_list_id(tasks);
    std::vector<server_task_result_ptr> results(task_ids.size());

    // Create a new HashMap instance
    jobject o_probabilities = env->NewObject(c_hash_map, cc_hash_map);
    if (o_probabilities == nullptr) {
        env->ThrowNew(c_llama_error, "Failed to create HashMap object.");
        return nullptr;
    }

    for (int i = 0; i < (int)task_ids.size(); i++) {
        server_task_result_ptr result = ctx_server->queue_results.recv(task_ids);
        if (result->is_error()) {
            auto response = result->to_json()["message"].get<std::string>();
            for (const int id_task : task_ids) {
                ctx_server->queue_results.remove_waiting_task_id(id_task);
            }
            env->ThrowNew(c_llama_error, response.c_str());
            return nullptr;
        }

        const auto out_res = result->to_json();

        if (result->is_stop()) {
            for (const int id_task : task_ids) {
                ctx_server->queue_results.remove_waiting_task_id(id_task);
            }
        }

        int index = out_res["index"].get<int>();
        float score = out_res["score"].get<float>();
        std::string tok_str = document_vector[index];
        jstring jtok_str = env->NewStringUTF(tok_str.c_str());

        jobject jprob = env->NewObject(c_float, cc_float, score);
        env->CallObjectMethod(o_probabilities, m_map_put, jtok_str, jprob);
        env->DeleteLocalRef(jtok_str);
        env->DeleteLocalRef(jprob);
    }
    jbyteArray jbytes = parse_jbytes(env, prompt);
    return env->NewObject(c_output, cc_output, jbytes, o_probabilities, true);
}

JNIEXPORT jstring JNICALL Java_de_kherud_llama_LlamaModel_applyTemplate(JNIEnv *env, jobject obj, jstring jparams) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    std::string c_params = parse_jstring(env, jparams);
    json data = json::parse(c_params);

    json templateData =
        oaicompat_completion_params_parse(data, ctx_server->params_base.use_jinja,
                                          ctx_server->params_base.reasoning_format, ctx_server->chat_templates.get());
    std::string tok_str = templateData.at("prompt");
    jstring jtok_str = env->NewStringUTF(tok_str.c_str());

    return jtok_str;
}

JNIEXPORT jintArray JNICALL Java_de_kherud_llama_LlamaModel_encode(JNIEnv *env, jobject obj, jstring jprompt) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    const std::string c_prompt = parse_jstring(env, jprompt);

    llama_tokens tokens = tokenize_mixed(ctx_server->vocab, c_prompt, false, true);
    jsize token_size = tokens.size(); // NOLINT(*-narrowing-conversions)

    jintArray java_tokens = env->NewIntArray(token_size);
    if (java_tokens == nullptr) {
        env->ThrowNew(c_error_oom, "could not allocate token memory");
        return nullptr;
    }

    env->SetIntArrayRegion(java_tokens, 0, token_size, reinterpret_cast<const jint *>(tokens.data()));

    return java_tokens;
}

JNIEXPORT jbyteArray JNICALL Java_de_kherud_llama_LlamaModel_decodeBytes(JNIEnv *env, jobject obj,
                                                                         jintArray java_tokens) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

    jsize length = env->GetArrayLength(java_tokens);
    jint *elements = env->GetIntArrayElements(java_tokens, nullptr);
    std::vector<llama_token> tokens(elements, elements + length);
    std::string text = tokens_to_str(ctx_server->ctx, tokens.cbegin(), tokens.cend());

    env->ReleaseIntArrayElements(java_tokens, elements, 0);

    return parse_jbytes(env, text);
}

JNIEXPORT void JNICALL Java_de_kherud_llama_LlamaModel_delete(JNIEnv *env, jobject obj) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)
    ctx_server->queue_tasks.terminate();
    // delete ctx_server;
}

JNIEXPORT void JNICALL Java_de_kherud_llama_LlamaModel_cancelCompletion(JNIEnv *env, jobject obj, jint id_task) {
    jlong server_handle = env->GetLongField(obj, f_model_pointer);
    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)
    std::unordered_set<int> id_tasks = {id_task};
    ctx_server->cancel_tasks(id_tasks);
    ctx_server->queue_results.remove_waiting_task_id(id_task);
}

JNIEXPORT void JNICALL Java_de_kherud_llama_LlamaModel_setLogger(JNIEnv *env, jclass clazz, jobject log_format,
                                                                 jobject jcallback) {
    if (o_log_callback != nullptr) {
        env->DeleteGlobalRef(o_log_callback);
    }

    log_json = env->IsSameObject(log_format, o_log_format_json);

    if (jcallback == nullptr) {
        log_callback = nullptr;
        llama_log_set(nullptr, nullptr);
    } else {
        o_log_callback = env->NewGlobalRef(jcallback);
        log_callback = [](enum ggml_log_level level, const char *text, void *user_data) {
            JNIEnv *env = get_jni_env();
            jstring message = env->NewStringUTF(text);
            jobject log_level = log_level_to_jobject(level);
            env->CallVoidMethod(o_log_callback, m_biconsumer_accept, log_level, message);
            env->DeleteLocalRef(message);
        };
        if (!log_json) {
            llama_log_set(log_callback_trampoline, nullptr);
        }
    }
}

JNIEXPORT jbyteArray JNICALL Java_de_kherud_llama_LlamaModel_jsonSchemaToGrammarBytes(JNIEnv *env, jclass clazz,
                                                                                      jstring j_schema) {
    const std::string c_schema = parse_jstring(env, j_schema);
    nlohmann::ordered_json c_schema_json = nlohmann::ordered_json::parse(c_schema);
    const std::string c_grammar = json_schema_to_grammar(c_schema_json);
    return parse_jbytes(env, c_grammar);
}