#include "RNLlamaJSI.h"
#include "JSIContext.h"
#include "ThreadPool.h"
#include "JSIUtils.h"
#include "JSIParams.h"
#include "JSIHelpers.h"
#include "JSISession.h"
#include "JSICompletion.h"
#include "JSIRequestManager.h"
#include "JSITaskManager.h"
#include "JSINativeHeaders.h"

#include <algorithm>
#include <atomic>
#include <map>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#if defined(__ANDROID__)
#include <android/log.h>
#include <cstring>
#endif

using namespace facebook;
using json = nlohmann::ordered_json;

// Consolidated logging function
enum class LogLevel { LOG_DEBUG, LOG_INFO, LOG_ERROR };

static void log(LogLevel level, const char* format, ...) {
    va_list args;
    va_start(args, format);

#if defined(__ANDROID__)
    int androidLevel = (level == LogLevel::LOG_DEBUG) ? ANDROID_LOG_DEBUG :
                      (level == LogLevel::LOG_INFO) ? ANDROID_LOG_INFO : ANDROID_LOG_ERROR;
    __android_log_vprint(androidLevel, "RNWhisperJSI", format, args);
#else
    char buffer[1024];
    vsnprintf(buffer, sizeof(buffer), format, args);
    const char* levelStr = (level == LogLevel::LOG_DEBUG) ? "DEBUG" :
                          (level == LogLevel::LOG_INFO) ? "INFO" : "ERROR";
    printf("RNWhisperJSI %s: %s\n", levelStr, buffer);
#endif

    va_end(args);
}

#define logInfo(format, ...) log(LogLevel::LOG_INFO, format, ##__VA_ARGS__)
#define logError(format, ...) log(LogLevel::LOG_ERROR, format, ##__VA_ARGS__)
#define logDebug(format, ...) log(LogLevel::LOG_DEBUG, format, ##__VA_ARGS__)

static std::once_flag backend_init_once;

#if defined(__ANDROID__)
static bool shouldExcludeHexagonDevice(lm_ggml_backend_dev_t dev) {
#if defined(LM_GGML_USE_HEXAGON)
    const char *dev_name = lm_ggml_backend_dev_name(dev);
    if (dev_name != nullptr && strncmp(dev_name, "HTP", 3) == 0) {
        return true;
    }
#else
    (void) dev;
#endif
    return false;
}

static std::vector<lm_ggml_backend_dev_t> getFilteredDefaultDevices() {
    std::vector<lm_ggml_backend_dev_t> rpc_servers;
    std::vector<lm_ggml_backend_dev_t> gpus;
    std::vector<lm_ggml_backend_dev_t> igpus;

    for (size_t i = 0; i < lm_ggml_backend_dev_count(); ++i) {
        lm_ggml_backend_dev_t dev = lm_ggml_backend_dev_get(i);
        if (shouldExcludeHexagonDevice(dev)) {
            continue;
        }

        switch (lm_ggml_backend_dev_type(dev)) {
            case LM_GGML_BACKEND_DEVICE_TYPE_CPU:
            case LM_GGML_BACKEND_DEVICE_TYPE_ACCEL:
            case LM_GGML_BACKEND_DEVICE_TYPE_META:
                break;
            case LM_GGML_BACKEND_DEVICE_TYPE_GPU: {
                lm_ggml_backend_reg_t reg = lm_ggml_backend_dev_backend_reg(dev);
                const char *reg_name = reg ? lm_ggml_backend_reg_name(reg) : nullptr;
                if (reg_name != nullptr && strcmp(reg_name, "RPC") == 0) {
                    rpc_servers.push_back(dev);
                } else {
                    lm_ggml_backend_dev_props props;
                    lm_ggml_backend_dev_get_props(dev, &props);
                    auto it = std::find_if(gpus.begin(), gpus.end(), [&props](lm_ggml_backend_dev_t other) {
                        lm_ggml_backend_dev_props other_props;
                        lm_ggml_backend_dev_get_props(other, &other_props);
                        return props.device_id != nullptr &&
                               other_props.device_id != nullptr &&
                               strcmp(props.device_id, other_props.device_id) == 0;
                    });

                    if (it == gpus.end()) {
                        gpus.push_back(dev);
                    }
                }
                break;
            }
            case LM_GGML_BACKEND_DEVICE_TYPE_IGPU:
                igpus.push_back(dev);
                break;
        }
    }

    std::vector<lm_ggml_backend_dev_t> devices;
    devices.insert(devices.end(), rpc_servers.begin(), rpc_servers.end());
    devices.insert(devices.end(), gpus.begin(), gpus.end());

    if (devices.empty()) {
        devices.insert(devices.end(), igpus.begin(), igpus.end());
    }

    if (!devices.empty()) {
        devices.push_back(nullptr);
    }

    return devices;
}
#endif

static std::string stripFileScheme(const std::string& path) {
    const std::string prefix = "file://";
    if (path.rfind(prefix, 0) == 0) {
        return path.substr(prefix.size());
    }
    return path;
}

namespace rnllama_jsi {
    static std::atomic<int64_t> g_context_limit(-1);
#if defined(__ANDROID__)
    static std::string g_android_loaded_library;
#endif
    static std::mutex g_log_mutex;
    static std::weak_ptr<react::CallInvoker> g_log_invoker;
    static std::shared_ptr<jsi::Function> g_log_handler;
    static std::shared_ptr<jsi::Runtime> g_log_runtime;

    struct ProgressCallbackData {
        std::shared_ptr<jsi::Function> callback;
        std::weak_ptr<react::CallInvoker> callInvoker;
        std::shared_ptr<jsi::Runtime> runtime;
        int contextId;
        std::atomic<int> lastProgress{0};
        int progressEvery = 1;
    };

    void setContextLimit(int64_t limit) {
        g_context_limit.store(limit);
    }

#if defined(__ANDROID__)
    void setAndroidLoadedLibrary(const std::string& name) {
        g_android_loaded_library = name;
    }
#endif

    static bool isContextLimitReached() {
        int64_t limit = g_context_limit.load();
        if (limit < 0) {
            return false;
        }
        return g_llamaContexts.size() >= static_cast<size_t>(limit);
    }

    static bool isContextBusy(rnllama::llama_rn_context* ctx) {
        if (ctx == nullptr) {
            return false;
        }

        if (ctx->completion && ctx->completion->is_predicting) {
            return true;
        }

        return ctx->slot_manager && ctx->slot_manager->has_pending_work();
    }

    static void throwIfContextBusy(rnllama::llama_rn_context* ctx) {
        if (isContextBusy(ctx)) {
            throw std::runtime_error("Context is busy");
        }
    }

    static void ensureBackendInitialized() {
        std::call_once(backend_init_once, []() {
            llama_backend_init();
        });
    }

    static void logToJsCallback(enum lm_ggml_log_level level, const char* text, void* /*data*/) {
        llama_log_callback_default(level, text, nullptr);

        std::shared_ptr<react::CallInvoker> invoker;
        std::shared_ptr<jsi::Function> handler;
        std::shared_ptr<jsi::Runtime> runtime;
        {
            std::lock_guard<std::mutex> lock(g_log_mutex);
            invoker = g_log_invoker.lock();
            handler = g_log_handler;
            runtime = g_log_runtime;
        }

        if (!invoker || !handler || !runtime) {
            return;
        }

        std::string levelStr = "info";
        switch (level) {
            case LM_GGML_LOG_LEVEL_ERROR: levelStr = "error"; break;
            case LM_GGML_LOG_LEVEL_WARN: levelStr = "warn"; break;
            case LM_GGML_LOG_LEVEL_INFO: levelStr = "info"; break;
            default: break;
        }

        std::string message = text ? text : "";

        invoker->invokeAsync([handler, levelStr, message, runtime]() {
            auto& rt = *runtime;
            handler->call(
                rt,
                jsi::String::createFromUtf8(rt, levelStr),
                jsi::String::createFromUtf8(rt, message)
            );
        });
    }

    // Helper: convert vector<string> to JSI array
    static jsi::Array toJsStringArray(jsi::Runtime& runtime, const std::vector<std::string>& values) {
        jsi::Array arr(runtime, values.size());
        for (size_t i = 0; i < values.size(); ++i) {
            arr.setValueAtIndex(runtime, i, jsi::String::createFromUtf8(runtime, values[i]));
        }
        return arr;
    }

    static bool isThinkingForcedOpen(const common_chat_params& chatParams) {
        if (!chatParams.supports_thinking || chatParams.thinking_start_tag.empty()) {
            return false;
        }

        const size_t lastStart = chatParams.generation_prompt.rfind(chatParams.thinking_start_tag);
        if (lastStart == std::string::npos) {
            return false;
        }

        for (const auto& endTag : chatParams.thinking_end_tags) {
            if (endTag.empty()) {
                continue;
            }
            const size_t lastEnd = chatParams.generation_prompt.rfind(endTag);
            if (lastEnd != std::string::npos && lastEnd >= lastStart) {
                return false;
            }
        }
        return true;
    }

    static jsi::Object createModelDetails(jsi::Runtime& runtime, rnllama::llama_rn_context* ctx) {
        jsi::Object model(runtime);

        char desc[1024];
        llama_model_desc(ctx->model, desc, sizeof(desc));
        model.setProperty(runtime, "desc", jsi::String::createFromUtf8(runtime, desc));
        model.setProperty(runtime, "size", (double)llama_model_size(ctx->model));
        model.setProperty(runtime, "nEmbd", (double)llama_model_n_embd(ctx->model));
        model.setProperty(runtime, "nParams", (double)llama_model_n_params(ctx->model));
        model.setProperty(runtime, "is_recurrent", llama_model_is_recurrent(ctx->model));
        model.setProperty(runtime, "is_hybrid", llama_model_is_hybrid(ctx->model));

        // Metadata
        jsi::Object metadata(runtime);
        int metaCount = llama_model_meta_count(ctx->model);
        for (int i = 0; i < metaCount; ++i) {
            char key[256];
            llama_model_meta_key_by_index(ctx->model, i, key, sizeof(key));
            char val[16384];
            llama_model_meta_val_str_by_index(ctx->model, i, val, sizeof(val));
            metadata.setProperty(runtime, key, jsi::String::createFromUtf8(runtime, val));
        }
        model.setProperty(runtime, "metadata", metadata);

        // Chat template capabilities
        jsi::Object chatTemplates(runtime);
        bool llamaChat = ctx->validateModelChatTemplate(false, nullptr);
        chatTemplates.setProperty(runtime, "llamaChat", llamaChat);

        jsi::Object jinja(runtime);
        bool jinjaDefault = ctx->validateModelChatTemplate(true, nullptr);
        jinja.setProperty(runtime, "default", jinjaDefault);

        jsi::Object defaultCaps(runtime);
        if (ctx->templates && common_chat_templates_has_variant(ctx->templates.get(), "")) {
            auto caps = common_chat_templates_get_caps(ctx->templates.get(), "");
            defaultCaps.setProperty(runtime, "tools", caps.supports_tools);
            defaultCaps.setProperty(runtime, "toolCalls", caps.supports_tool_calls);
            defaultCaps.setProperty(runtime, "parallelToolCalls", caps.supports_parallel_tool_calls);
            defaultCaps.setProperty(runtime, "systemRole", caps.supports_system_role);
        } else {
            defaultCaps.setProperty(runtime, "tools", false);
            defaultCaps.setProperty(runtime, "toolCalls", false);
            defaultCaps.setProperty(runtime, "parallelToolCalls", false);
            defaultCaps.setProperty(runtime, "systemRole", false);
        }
        jinja.setProperty(runtime, "defaultCaps", defaultCaps);

        bool toolUseSupported = ctx->validateModelChatTemplate(true, "tool_use");
        jinja.setProperty(runtime, "toolUse", toolUseSupported);
        if (ctx->templates && common_chat_templates_has_variant(ctx->templates.get(), "tool_use")) {
            auto caps = common_chat_templates_get_caps(ctx->templates.get(), "tool_use");
            jsi::Object toolUseCaps(runtime);
            toolUseCaps.setProperty(runtime, "tools", caps.supports_tools);
            toolUseCaps.setProperty(runtime, "toolCalls", caps.supports_tool_calls);
            toolUseCaps.setProperty(runtime, "parallelToolCalls", caps.supports_parallel_tool_calls);
            toolUseCaps.setProperty(runtime, "systemRole", caps.supports_system_role);
            jinja.setProperty(runtime, "toolUseCaps", toolUseCaps);
        }

        chatTemplates.setProperty(runtime, "jinja", jinja);
        model.setProperty(runtime, "chatTemplates", chatTemplates);

        // Deprecated flag maintained for compatibility
        model.setProperty(runtime, "isChatTemplateSupported", llamaChat);

        return model;
    }

    static std::vector<lm_ggml_backend_dev_t> buildDeviceOverrides(
        const std::vector<std::string>& requestedDevices,
        bool skipGpuDevices,
        bool& anyGpuAvailable
    ) {
        std::vector<lm_ggml_backend_dev_t> selected;
        anyGpuAvailable = false;

        const size_t devCount = lm_ggml_backend_dev_count();
        for (size_t i = 0; i < devCount; ++i) {
            lm_ggml_backend_dev_t dev = lm_ggml_backend_dev_get(i);
            const auto type = lm_ggml_backend_dev_type(dev);
#if TARGET_OS_SIMULATOR
            if (type == LM_GGML_BACKEND_DEVICE_TYPE_ACCEL) {
                continue;
            }
#endif
            const bool isGpuType = type == LM_GGML_BACKEND_DEVICE_TYPE_GPU || type == LM_GGML_BACKEND_DEVICE_TYPE_IGPU;
            if (isGpuType) {
                anyGpuAvailable = true;
            }
            if (skipGpuDevices && isGpuType) {
                continue;
            }

            if (!requestedDevices.empty()) {
                const char* name = lm_ggml_backend_dev_name(dev);
                std::string nameStr = name ? name : "";
                auto it = std::find(requestedDevices.begin(), requestedDevices.end(), nameStr);
                if (it == requestedDevices.end()) {
                    continue;
                }
            }

            selected.push_back(dev);
        }

        if (!selected.empty()) {
            selected.push_back(nullptr);
        }

        return selected;
    }

    static bool isGpuDeviceType(enum lm_ggml_backend_dev_type type) {
        return type == LM_GGML_BACKEND_DEVICE_TYPE_GPU || type == LM_GGML_BACKEND_DEVICE_TYPE_IGPU;
    }

    static bool hasGpuBackendDevice() {
        const size_t devCount = lm_ggml_backend_dev_count();
        for (size_t i = 0; i < devCount; ++i) {
            auto dev = lm_ggml_backend_dev_get(i);
            if (isGpuDeviceType(lm_ggml_backend_dev_type(dev))) {
                return true;
            }
        }
        return false;
    }

    static void configureBackendDevices(
        common_params& cparams,
        const std::vector<std::string>& requestedDevices,
        bool devicesProvided,
        bool skipGpuDevices,
        bool& anyGpuAvailable
    ) {
        anyGpuAvailable = false;
        std::vector<lm_ggml_backend_dev_t> overrideDevices;

        if (devicesProvided) {
            overrideDevices = buildDeviceOverrides(requestedDevices, skipGpuDevices, anyGpuAvailable);
            if (!overrideDevices.empty()) {
                cparams.devices = overrideDevices;
            }
        }

        if (overrideDevices.empty() && !skipGpuDevices) {
#if defined(__ANDROID__)
            auto defaultDevices = getFilteredDefaultDevices();
            if (!defaultDevices.empty()) {
                cparams.devices = defaultDevices;
                for (auto dev : defaultDevices) {
                    if (dev == nullptr) continue;
                    if (isGpuDeviceType(lm_ggml_backend_dev_type(dev))) {
                        anyGpuAvailable = true;
                        break;
                    }
                }
            }
#endif
        }

        // Track backend availability when no explicit override was applied.
        if (overrideDevices.empty() && !anyGpuAvailable) {
            anyGpuAvailable = hasGpuBackendDevice();
        }
    }

    void addContext(int contextId, long contextPtr) {
        g_llamaContexts.add(contextId, contextPtr);
    }

    void removeContext(int contextId) {
        g_llamaContexts.remove(contextId);
    }

    rnllama::llama_rn_context* getContextOrThrow(int contextId) {
        long ctxPtr = g_llamaContexts.get(contextId);
        if (!ctxPtr) {
            throw std::runtime_error("Context not found");
        }
        return reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
    }

    void installJSIBindings(
        jsi::Runtime& runtime,
        std::shared_ptr<react::CallInvoker> callInvoker
    ) {
        TaskManager::getInstance().reset();
        auto initContext = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaInitContext"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Object params = arguments[1].asObject(runtime);
                bool isModelAsset = getPropertyAsBool(runtime, params, "is_model_asset", false);
                bool isModelDraftAsset = getPropertyAsBool(runtime, params, "is_model_draft_asset", false);

                bool useProgressCallback = getPropertyAsBool(runtime, params, "use_progress_callback", false);
                int progressCallbackEvery = getPropertyAsInt(runtime, params, "progress_callback_every", 1);
                std::shared_ptr<ProgressCallbackData> progressData;
                if (count > 2 && arguments[2].isObject() && arguments[2].asObject(runtime).isFunction(runtime)) {
                    useProgressCallback = true;
                    progressData = std::make_shared<ProgressCallbackData>();
                    progressData->callback = makeJsiFunction(runtime, arguments[2], callInvoker);
                    progressData->callInvoker = callInvoker;
                    progressData->runtime = std::shared_ptr<jsi::Runtime>(&runtime, [](jsi::Runtime*){});
                    progressData->contextId = contextId;
                    progressData->progressEvery = std::max(1, progressCallbackEvery);
                    progressData->lastProgress.store(0);
                } else if (useProgressCallback) {
                    // Progress requested but no callback provided
                    useProgressCallback = false;
                }

                common_params cparams;
                parseCommonParams(runtime, params, cparams);

#if defined(__APPLE__)
                if (isModelAsset) {
                    cparams.model.path = resolveIosModelPath(cparams.model.path, true);
                }
                if (isModelDraftAsset && !cparams.speculative.draft.mparams.path.empty()) {
                    cparams.speculative.draft.mparams.path =
                        resolveIosModelPath(cparams.speculative.draft.mparams.path, true);
                }
#endif

                bool skipGpuDevices = getPropertyAsBool(runtime, params, "no_gpu_devices", false);
                if (skipGpuDevices) {
                    cparams.n_gpu_layers = 0;
                }

                std::vector<std::string> requestedDevices;
                bool devicesProvided = false;
                if (params.hasProperty(runtime, "devices") && params.getProperty(runtime, "devices").isObject()) {
                    jsi::Array devicesArr = params.getProperty(runtime, "devices").asObject(runtime).asArray(runtime);
                    if (devicesArr.size(runtime) > 0) {
                        devicesProvided = true;
                        for (size_t i = 0; i < devicesArr.size(runtime); ++i) {
                            auto val = devicesArr.getValueAtIndex(runtime, i);
                            if (val.isString()) {
                                requestedDevices.push_back(val.asString(runtime).utf8(runtime));
                            }
                        }
                    }
                }

                int stateCacheBudgetMb =
                    getPropertyAsInt(runtime, params, "state_cache_budget_mb", 160);
                int stateCacheMaxCheckpoints =
                    getPropertyAsInt(runtime, params, "state_cache_max_checkpoints", 8);

                return createPromiseTask(runtime, callInvoker, [
                    contextId,
                    cparams,
                    skipGpuDevices,
                    requestedDevices,
                    devicesProvided,
                    useProgressCallback,
                    progressData,
                    stateCacheBudgetMb,
                    stateCacheMaxCheckpoints
                ]() mutable -> PromiseResultGenerator {
                    if (isContextLimitReached()) {
                        throw std::runtime_error("Context limit reached");
                    }

                    ensureBackendInitialized();

#if defined(__APPLE__)
                    auto metalAvailability = getMetalAvailability(skipGpuDevices);
                    std::string appleGpuReason = metalAvailability.available ? "" : metalAvailability.reason;
                    if (!metalAvailability.available && !skipGpuDevices) {
                        skipGpuDevices = true;
                        cparams.n_gpu_layers = 0;
                    }
#endif

                    bool anyGpuAvailable = false;
                    configureBackendDevices(
                        cparams,
                        requestedDevices,
                        devicesProvided,
                        skipGpuDevices,
                        anyGpuAvailable
                    );

                    if (useProgressCallback && progressData && progressData->callback) {
                        cparams.progress_callback = [](float progress, void * user_data) {
                            auto *data = static_cast<ProgressCallbackData *>(user_data);
                            if (!data) {
                                return true;
                            }

                            int percentage = (int) (progress * 100.0f);
                            int last = data->lastProgress.load();
                            if (percentage < 100 && percentage - last < data->progressEvery) {
                                return true;
                            }
                            if (percentage <= last) {
                                return true;
                            }

                            data->lastProgress.store(percentage);

                            auto invoker = data->callInvoker.lock();
                            auto cb = data->callback;
                            auto runtime = data->runtime;
                            if (invoker && cb && runtime) {
                                invoker->invokeAsync([cb, percentage, runtime]() {
                                    auto& rt = *runtime;
                                    cb->call(rt, jsi::Value((double) percentage));
                                });
                            }

                            return true;
                        };
                        cparams.progress_callback_user_data = progressData.get();
                    }

                    auto ctx = new rnllama::llama_rn_context();
                    // Prompt state cache tuning (multi-turn KV reuse on
                    // recurrent/hybrid/SWA models). Budget in MiB; 0 disables it.
                    {
                        ctx->state_cache_budget_bytes =
                            stateCacheBudgetMb > 0 ? (size_t) stateCacheBudgetMb * 1024 * 1024 : 0;
                        ctx->state_cache_max_checkpoints = stateCacheMaxCheckpoints;
                    }
                    if (ctx->loadModel(cparams)) {
                         ctx->attachThreadpoolsIfAvailable();

                         if (ctx->params.embedding && llama_model_has_encoder(ctx->model) && llama_model_has_decoder(ctx->model)) {
                             delete ctx;
                             throw std::runtime_error("Embedding is not supported in encoder-decoder models");
                         }

                         std::vector<std::string> usedDevices;
                         bool gpuEnabled = false;
                         if (ctx->llama_init->model() != nullptr) {
                             for (const auto & dev_info : ctx->llama_init->model()->devices) {
                                 auto dev = dev_info.dev;
                                 if (dev == nullptr) continue;
                                 const char* used_name = lm_ggml_backend_dev_name(dev);
                                 if (used_name != nullptr) {
                                     usedDevices.push_back(used_name);
                                 }
                                 if (isGpuDeviceType(lm_ggml_backend_dev_type(dev))) {
                                     gpuEnabled = true;
                                 }
                             }
                         }

                         std::string reasonNoGPU;
#if defined(__APPLE__)
                         const std::string platformReason = appleGpuReason;
#endif
                         if (!gpuEnabled) {
#if defined(__APPLE__)
                             if (!platformReason.empty()) {
                                 reasonNoGPU = platformReason;
                             } else
#endif
                             if (skipGpuDevices) {
                                 reasonNoGPU = "GPU devices disabled by user";
                             } else if (anyGpuAvailable) {
                                 reasonNoGPU = "GPU backend is available but was not selected";
                             } else {
                                 reasonNoGPU = "GPU backend is not available";
                             }
                         }

                         addContext(contextId, (long)ctx);

                         std::string system_info = common_params_get_system_info(ctx->params);

                         return [gpuEnabled, reasonNoGPU, system_info, usedDevices, contextId](jsi::Runtime& rt) {
                             jsi::Object result(rt);
                             result.setProperty(rt, "gpu", gpuEnabled);
                             result.setProperty(rt, "reasonNoGPU", jsi::String::createFromUtf8(rt, reasonNoGPU));
                             result.setProperty(rt, "systemInfo", jsi::String::createFromUtf8(rt, system_info));

                             // Model metadata and chat template capabilities
                             long ctxPtr = g_llamaContexts.get(contextId);
                             if (ctxPtr) {
                                 auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                                 result.setProperty(rt, "model", createModelDetails(rt, ctx));
                             }

                             // Maintain shape expected by TypeScript
                             result.setProperty(rt, "devices", toJsStringArray(rt, usedDevices));
                             std::string androidLibName = "";
                             #if defined(__ANDROID__)
                             androidLibName = g_android_loaded_library;
                             #endif
                             result.setProperty(rt, "androidLib", jsi::String::createFromUtf8(rt, androidLibName));
                             return result;
                         };
                    } else {
                        delete ctx;
                        throw std::runtime_error("Failed to load model");
                    }
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaInitContext", initContext);

        // ... (modelInfo, getBackendDevicesInfo, loadSession, saveSession, tokenize, detokenize, getFormattedChat from previous)
        auto modelInfo = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaModelInfo"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                std::string path = arguments[0].asString(runtime).utf8(runtime);
                std::vector<std::string> skip;
                if (count > 1 && arguments[1].isObject()) {
                    jsi::Array skipArr = arguments[1].asObject(runtime).asArray(runtime);
                    for (size_t i = 0; i < skipArr.size(runtime); i++) {
                        skip.push_back(skipArr.getValueAtIndex(runtime, i).asString(runtime).utf8(runtime));
                    }
                }

                return createPromiseTask(runtime, callInvoker, [path, skip]() -> PromiseResultGenerator {
                    return [path, skip](jsi::Runtime& rt) {
                        return createModelInfo(rt, path, skip);
                    };
                }, -1, false);
            }
        );
        runtime.global().setProperty(runtime, "llamaModelInfo", modelInfo);

        auto getBackendDevicesInfo = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetBackendDevicesInfo"),
            0,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                 return createPromiseTask(runtime, callInvoker, [callInvoker]() -> PromiseResultGenerator {
                     ensureBackendInitialized();

                     std::string info = rnllama::get_backend_devices_info();

                     return [info](jsi::Runtime& rt) {
                         return jsi::String::createFromUtf8(rt, info);
                     };
                 }, -1, false);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetBackendDevicesInfo", getBackendDevicesInfo);

        auto loadSession = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaLoadSession"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string path = arguments[1].asString(runtime).utf8(runtime);

                return createPromiseTask(runtime, callInvoker, [contextId, path]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    return [contextId, path](jsi::Runtime& rt) {
                        long ctxPtr = g_llamaContexts.get(contextId);
                        if (!ctxPtr) {
                            throw std::runtime_error("Context was released");
                        }
                        auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                        return rnllama_jsi::loadSession(rt, ctx, path);
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaLoadSession", loadSession);

        auto saveSession = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaSaveSession"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string path = arguments[1].asString(runtime).utf8(runtime);
                int size = (int)arguments[2].asNumber();

                return createPromiseTask(runtime, callInvoker, [contextId, path, size]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    int tokens_saved = rnllama_jsi::saveSession(ctx, path, size);
                    return [tokens_saved](jsi::Runtime& rt) {
                        return jsi::Value(tokens_saved);
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaSaveSession", saveSession);

        auto tokenize = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaTokenize"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string text = arguments[1].asString(runtime).utf8(runtime);
                std::vector<std::string> mediaPaths;
                if (count > 2 && arguments[2].isObject()) {
                    jsi::Array paths = arguments[2].asObject(runtime).asArray(runtime);
                    for (size_t i = 0; i < paths.size(runtime); ++i) {
                         mediaPaths.push_back(paths.getValueAtIndex(runtime, i).asString(runtime).utf8(runtime));
                    }
                }

                return createPromiseTask(runtime, callInvoker, [contextId, text, mediaPaths]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    auto result = ctx->tokenize(text, mediaPaths);
                    return [result](jsi::Runtime& rt) {
                        return createTokenizeResult(rt, result);
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaTokenize", tokenize);

        auto detokenize = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaDetokenize"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::vector<llama_token> tokens;
                jsi::Array tokensArr = arguments[1].asObject(runtime).asArray(runtime);
                for (size_t i = 0; i < tokensArr.size(runtime); ++i) {
                    tokens.push_back((llama_token)tokensArr.getValueAtIndex(runtime, i).asNumber());
                }

                return createPromiseTask(runtime, callInvoker, [contextId, tokens]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    std::string text = rnllama::tokens_to_str(ctx->ctx, tokens.cbegin(), tokens.cend());
                    return [text](jsi::Runtime& rt) {
                        return jsi::String::createFromUtf8(rt, text);
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaDetokenize", detokenize);

        auto getFormattedChat = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetFormattedChat"),
            4,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                 int contextId = (int)arguments[0].asNumber();
                 std::string messages = arguments[1].asString(runtime).utf8(runtime);
                 std::string chatTemplate = "";
                 if (count > 2 && arguments[2].isString()) {
                     chatTemplate = arguments[2].asString(runtime).utf8(runtime);
                 }

                 std::string jsonSchema = "";
                 std::string tools = "";
                 bool parallelToolCalls = false;
                 std::string toolChoice = "";
                 bool enableThinking = false;
                 std::string reasoningFormat = "none";
                 bool addGenerationPrompt = true;
                 std::string nowStr = "";
                 std::map<std::string, std::string> chatTemplateKwargs;
                 bool useJinja = false;
                 bool forcePureContent = false;

                 if (count > 3 && arguments[3].isObject()) {
                     jsi::Object params = arguments[3].asObject(runtime);
                     useJinja = getPropertyAsBool(runtime, params, "jinja", false);

                     if (useJinja) {
                         jsonSchema = getPropertyAsString(runtime, params, "json_schema");
                         tools = getPropertyAsString(runtime, params, "tools");
                         parallelToolCalls = getPropertyAsBool(runtime, params, "parallel_tool_calls", false);
                         toolChoice = getPropertyAsString(runtime, params, "tool_choice");
                         enableThinking = getPropertyAsBool(runtime, params, "enable_thinking", false);
                         reasoningFormat = getPropertyAsString(runtime, params, "reasoning_format", "none");
                         addGenerationPrompt = getPropertyAsBool(runtime, params, "add_generation_prompt", true);
                         nowStr = getPropertyAsString(runtime, params, "now");
                         forcePureContent = getPropertyAsBool(runtime, params, "force_pure_content", false);

                         std::string kwargsStr = getPropertyAsString(runtime, params, "chat_template_kwargs");
                          if (!kwargsStr.empty()) {
                              try {
                                  auto kwargs_json = json::parse(kwargsStr);
                                  for (auto& [key, value] : kwargs_json.items()) {
                                      if (value.is_string()) {
                                          chatTemplateKwargs[key] = value.get<std::string>();
                                      }
                                  }
                              } catch (...) { }
                          }
                     }
                 }

                 return createPromiseTask(runtime, callInvoker, [contextId, messages, chatTemplate, jsonSchema, tools, parallelToolCalls, toolChoice, enableThinking, reasoningFormat, addGenerationPrompt, nowStr, chatTemplateKwargs, useJinja, forcePureContent]() -> PromiseResultGenerator {
                      auto ctx = getContextOrThrow(contextId);
                      if (useJinja) {
                          auto chatParams = ctx->getFormattedChatWithJinja(
                               messages, chatTemplate, jsonSchema, tools, parallelToolCalls,
                               toolChoice, enableThinking, reasoningFormat, addGenerationPrompt, nowStr, chatTemplateKwargs, forcePureContent
                          );

                          return [chatParams](jsi::Runtime& rt) {
                              jsi::Object result(rt);
                              result.setProperty(rt, "prompt", jsi::String::createFromUtf8(rt, chatParams.prompt));
                              result.setProperty(rt, "chat_format", (int)chatParams.format);
                              result.setProperty(rt, "grammar", jsi::String::createFromUtf8(rt, chatParams.grammar));
                              result.setProperty(rt, "grammar_lazy", chatParams.grammar_lazy);
                              result.setProperty(rt, "generation_prompt", jsi::String::createFromUtf8(rt, chatParams.generation_prompt));
                              result.setProperty(rt, "thinking_forced_open", isThinkingForcedOpen(chatParams));
                              if (!chatParams.thinking_start_tag.empty()) {
                                  result.setProperty(rt, "thinking_start_tag", jsi::String::createFromUtf8(rt, chatParams.thinking_start_tag));
                              }
                              if (!chatParams.thinking_end_tags.empty()) {
                                  result.setProperty(rt, "thinking_end_tag", jsi::String::createFromUtf8(rt, chatParams.thinking_end_tags.front()));
                              }

                              // Preserve the same shape as legacy native bridge
                              result.setProperty(rt, "type", jsi::String::createFromUtf8(rt, "jinja"));

                              jsi::Array preserved(rt, chatParams.preserved_tokens.size());
                              for (size_t i = 0; i < chatParams.preserved_tokens.size(); i++) {
                                  preserved.setValueAtIndex(rt, i, jsi::String::createFromUtf8(rt, chatParams.preserved_tokens[i]));
                              }
                              result.setProperty(rt, "preserved_tokens", preserved);

                              jsi::Array additionalStops(rt, chatParams.additional_stops.size());
                              for (size_t i = 0; i < chatParams.additional_stops.size(); i++) {
                                  additionalStops.setValueAtIndex(rt, i, jsi::String::createFromUtf8(rt, chatParams.additional_stops[i]));
                              }
                              result.setProperty(rt, "additional_stops", additionalStops);

                              jsi::Array triggers = jsi::Array(rt, chatParams.grammar_triggers.size());
                              for (size_t i = 0; i < chatParams.grammar_triggers.size(); i++) {
                                  jsi::Object trigger(rt);
                                  trigger.setProperty(rt, "type", (int)chatParams.grammar_triggers[i].type);
                                  trigger.setProperty(rt, "value", jsi::String::createFromUtf8(rt, chatParams.grammar_triggers[i].value));
                                  trigger.setProperty(rt, "token", (int)chatParams.grammar_triggers[i].token);
                                  triggers.setValueAtIndex(rt, i, trigger);
                              }
                              result.setProperty(rt, "grammar_triggers", triggers);

                              // Return the PEG parser string for COMMON_CHAT_FORMAT_PEG_* formats
                              if (!chatParams.parser.empty()) {
                                  result.setProperty(rt, "chat_parser", jsi::String::createFromUtf8(rt, chatParams.parser));
                              }

                              return result;
                          };
                      } else {
                          std::string prompt = ctx->getFormattedChat(messages, chatTemplate);
                          return [prompt](jsi::Runtime& rt) {
                              return jsi::String::createFromUtf8(rt, prompt);
                          };
                      }
                 }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetFormattedChat", getFormattedChat);

        auto embedding = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaEmbedding"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string text = arguments[1].asString(runtime).utf8(runtime);
                jsi::Object params = arguments[2].asObject(runtime);

                int embd_normalize = 0;
                bool has_embd_normalize = false;
                if (params.hasProperty(runtime, "embd_normalize")) {
                    embd_normalize = getPropertyAsInt(runtime, params, "embd_normalize", 2);
                    has_embd_normalize = true;
                }

                return createPromiseTask(runtime, callInvoker, [contextId, text, embd_normalize, has_embd_normalize]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);

                    if (!ctx->completion) throw std::runtime_error("Completion not initialized");
                    if (ctx->params.embedding != true) throw std::runtime_error("Embedding is not enabled");
                    throwIfContextBusy(ctx);

                    common_params embdParams = ctx->params;
                    embdParams.embedding = true;
                    embdParams.embd_normalize = has_embd_normalize ? embd_normalize : ctx->params.embd_normalize;

                    ctx->params.prompt = text;
                    ctx->params.n_predict = 0;

                    std::vector<float> result = ctx->completion->embedding(embdParams);

                    return [result](jsi::Runtime& rt) {
                        jsi::Object resultDict(rt);
                        jsi::Array embeddingResult(rt, result.size());
                        for (size_t i = 0; i < result.size(); i++) {
                            embeddingResult.setValueAtIndex(rt, i, (double)result[i]);
                        }
                        resultDict.setProperty(rt, "embedding", embeddingResult);
                        return resultDict;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaEmbedding", embedding);

        auto rerank = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaRerank"),
            4,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string query = arguments[1].asString(runtime).utf8(runtime);
                jsi::Array documentsArr = arguments[2].asObject(runtime).asArray(runtime);
                std::vector<std::string> documents;
                for (size_t i = 0; i < documentsArr.size(runtime); i++) {
                    documents.push_back(documentsArr.getValueAtIndex(runtime, i).asString(runtime).utf8(runtime));
                }
                // params argument ignored for now as per iOS implementation logic (only checks context state)

                return createPromiseTask(runtime, callInvoker, [contextId, query, documents]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);

                    if (!ctx->completion) throw std::runtime_error("Completion not initialized");
                    if (ctx->params.embedding != true) throw std::runtime_error("Embedding is not enabled");
                    throwIfContextBusy(ctx);

                    std::vector<float> scores = ctx->completion->rerank(query, documents);

                    return [scores](jsi::Runtime& rt) {
                        jsi::Array result(rt, scores.size());
                        for (size_t i = 0; i < scores.size(); i++) {
                            jsi::Object item(rt);
                            item.setProperty(rt, "score", (double)scores[i]);
                            item.setProperty(rt, "index", (int)i);
                            result.setValueAtIndex(rt, i, item);
                        }
                        return result;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaRerank", rerank);

        auto bench = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaBench"),
            5,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                int pp = (int)arguments[1].asNumber();
                int tg = (int)arguments[2].asNumber();
                int pl = (int)arguments[3].asNumber();
                int nr = (int)arguments[4].asNumber();

                return createPromiseTask(runtime, callInvoker, [contextId, pp, tg, pl, nr]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->completion) return [](jsi::Runtime& rt) { return jsi::String::createFromUtf8(rt, ""); };

                    std::string res = ctx->completion->bench(pp, tg, pl, nr);

                    return [res](jsi::Runtime& rt) {
                        return jsi::String::createFromUtf8(rt, res);
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaBench", bench);

        auto completion = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaCompletion"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Object params = arguments[1].asObject(runtime);
                std::shared_ptr<jsi::Function> onToken;

                if (count > 2 && arguments[2].isObject() && arguments[2].asObject(runtime).isFunction(runtime)) {
                    onToken = makeJsiFunction(runtime, arguments[2], callInvoker);
                }

                bool emitPartial = getPropertyAsBool(runtime, params, "emit_partial_completion", false);

                auto ctx = getContextOrThrow(contextId);
                throwIfContextBusy(ctx);
                ctx->completion->rewind();

                parseCompletionParams(runtime, params, ctx);

                std::vector<std::string> mediaPaths;
                if (params.hasProperty(runtime, "media_paths")) {
                    jsi::Array paths = params.getProperty(runtime, "media_paths").asObject(runtime).asArray(runtime);
                    for (size_t i = 0; i < paths.size(runtime); i++) {
                        mediaPaths.push_back(paths.getValueAtIndex(runtime, i).asString(runtime).utf8(runtime));
                    }
                }

                int chat_format = getPropertyAsInt(runtime, params, "chat_format", 0);
                std::string reasoningFormatStr = getPropertyAsString(runtime, params, "reasoning_format", "none");
                common_reasoning_format reasoning_format = common_reasoning_format_from_name(reasoningFormatStr);
                std::string generation_prompt = getPropertyAsString(runtime, params, "generation_prompt");
                std::string chat_parser = getPropertyAsString(runtime, params, "chat_parser");
                std::string prefill_text = getPropertyAsString(runtime, params, "prefill_text");

                return createPromiseTask(runtime, callInvoker, [runtimePtr = std::shared_ptr<jsi::Runtime>(&runtime, [](jsi::Runtime*){}), contextId, onToken, emitPartial, mediaPaths, chat_format, reasoning_format, generation_prompt, chat_parser, prefill_text, callInvoker]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);

                    if (ctx->completion == nullptr) {
                        throw std::runtime_error("Completion not initialized");
                    }
                    throwIfContextBusy(ctx);

                    if (!mediaPaths.empty() && ctx->completion->shouldUseMTP()) {
                        throw std::runtime_error("MTP speculative decoding currently supports text-only completion");
                    }

                    // NOTE: no rewind() here — it already ran on the JS thread
                    // BEFORE parseCompletionParams. rewind() resets
                    // sampling.grammar / antiprompt, so calling it again at this
                    // point would wipe the grammar and stop sequences this
                    // completion's params just configured (that regression broke
                    // TTS grammar-forced output).
                    if (!ctx->completion->initSampling()) {
                        throw std::runtime_error("Failed to initialize sampling");
                    }

                    ctx->completion->prefill_text = rnllama::utf8_sanitize(prefill_text);
                    ctx->completion->beginCompletion(chat_format, reasoning_format, generation_prompt, chat_parser);

                    try {
                        if (!mediaPaths.empty() && !ctx->isMultimodalEnabled()) {
                            throw std::runtime_error("Multimodal support not enabled. Call initMultimodal first.");
                        }
                        ctx->completion->loadPrompt(mediaPaths);
                    } catch (const std::exception &e) {
                        ctx->completion->endCompletion();
                        throw std::runtime_error(e.what());
                    }

                    if (ctx->completion->context_full) {
                        ctx->completion->endCompletion();
                        throw std::runtime_error("Context is full");
                    }

                    size_t sent_count = 0;

                    while (ctx->completion->has_next_token && !ctx->completion->is_interrupted) {
                        const rnllama::completion_token_output token_with_probs = ctx->completion->doCompletion();
                        if (token_with_probs.tok == -1 || ctx->completion->incomplete) {
                            continue;
                        }

                        const std::string token_text = common_token_to_piece(ctx->ctx, token_with_probs.tok);
                        size_t pos = std::min(sent_count, ctx->completion->generated_text.size());
                        const std::string str_test = ctx->completion->generated_text.substr(pos);

                        bool is_stop_full = false;
                        size_t stop_pos = ctx->completion->findStoppingStrings(str_test, token_text.size(), rnllama::STOP_FULL);
                        if (stop_pos != std::string::npos) {
                            is_stop_full = true;
                            ctx->completion->generated_text.erase(
                                ctx->completion->generated_text.begin() + pos + stop_pos,
                                ctx->completion->generated_text.end());
                            pos = std::min(sent_count, ctx->completion->generated_text.size());
                        } else {
                             stop_pos = ctx->completion->findStoppingStrings(str_test, token_text.size(), rnllama::STOP_PARTIAL);
                        }

                        if (stop_pos == std::string::npos || (!ctx->completion->has_next_token && !is_stop_full && stop_pos > 0)) {
                            const std::string to_send = ctx->completion->generated_text.substr(pos, std::string::npos);
                            sent_count += to_send.size();

                            if (emitPartial && onToken) {
                                rnllama::completion_token_output output_copy = token_with_probs;
                                output_copy.text = to_send;

                                rnllama::completion_chat_output partial_output;
                                bool has_partial_output = false;
                                try {
                                    partial_output = ctx->completion->parseChatOutput(true);
                                    has_partial_output = true;
                                } catch (...) {
                                    // ignore parse errors for partial output
                                }

                                auto runtime = runtimePtr;
                                if (runtime) {
                                    callInvoker->invokeAsync([onToken, output_copy, contextId, partial_output, has_partial_output, runtime]() {
                                        // Check if context is still valid (may have been released during async callback)
                                        long ctxPtr = g_llamaContexts.get(contextId);
                                        if (!ctxPtr) {
                                            // Context was released, skip token callback
                                            return;
                                        }
                                        auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                                        auto& rt = *runtime;
                                        jsi::Object res = createTokenResult(rt, ctx, output_copy);
                                        if (has_partial_output) {
                                            setChatOutputFields(rt, res, partial_output);
                                        }
                                        onToken->call(rt, res);
                                    });
                                }
                            }
                        }
                    }

                    common_perf_print(ctx->ctx, ctx->completion->ctx_sampling);
                    ctx->completion->endCompletion();

                    return [contextId](jsi::Runtime& rt) -> jsi::Value {
                        // Check if context is still valid (may have been released during async callback)
                        long ctxPtr = g_llamaContexts.get(contextId);
                        if (!ctxPtr) {
                            // Context was released, return minimal interrupted result
                            jsi::Object res(rt);
                            res.setProperty(rt, "text", jsi::String::createFromUtf8(rt, ""));
                            res.setProperty(rt, "interrupted", true);
                            res.setProperty(rt, "context_released", true);
                            return jsi::Value(std::move(res));
                        }
                        auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                        return jsi::Value(std::move(createCompletionResult(rt, ctx)));
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaCompletion", completion);

        auto stopCompletion = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaStopCompletion"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                auto ctx = getContextOrThrow(contextId);
                if (ctx->completion) {
                    ctx->completion->is_interrupted = true;
                }
                return jsi::Value::undefined();
            }
        );
        runtime.global().setProperty(runtime, "llamaStopCompletion", stopCompletion);

        auto toggleNativeLog = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaToggleNativeLog"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                bool enabled = count > 0 && arguments[0].isBool() ? arguments[0].getBool() : false;
                std::shared_ptr<jsi::Function> onLog;
                if (enabled && count > 1 && arguments[1].isObject() && arguments[1].asObject(runtime).isFunction(runtime)) {
                    onLog = makeJsiFunction(runtime, arguments[1], callInvoker);
                }

                return createPromiseTask(runtime, callInvoker, [enabled, onLog, callInvoker, runtimePtr = std::shared_ptr<jsi::Runtime>(&runtime, [](jsi::Runtime*){})]() -> PromiseResultGenerator {
                    if (enabled && onLog) {
                        {
                            std::lock_guard<std::mutex> lock(g_log_mutex);
                            g_log_handler = onLog;
                            g_log_invoker = callInvoker;
                            g_log_runtime = runtimePtr;
                        }
                        llama_log_set(logToJsCallback, nullptr);
                    } else {
                        {
                            std::lock_guard<std::mutex> lock(g_log_mutex);
                            g_log_handler.reset();
                            g_log_invoker.reset();
                            g_log_runtime.reset();
                        }
                        llama_log_set(llama_log_callback_default, nullptr);
                    }
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, -1, false);
            }
        );
        runtime.global().setProperty(runtime, "llamaToggleNativeLog", toggleNativeLog);

        auto enableParallelMode = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaEnableParallelMode"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Object params = arguments[1].asObject(runtime);

                bool enabled = getPropertyAsBool(runtime, params, "enabled", true);
                int nParallel = getPropertyAsInt(runtime, params, "n_parallel", 2);
                int nBatch = getPropertyAsInt(runtime, params, "n_batch", 512);

                return createPromiseTask(runtime, callInvoker, [contextId, enabled, nParallel, nBatch]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (enabled) {
                        ctx->enableParallelMode(nParallel, nBatch);
                        if (ctx->slot_manager) {
                            ctx->slot_manager->start_processing_loop();
                        }
                    } else {
                        if (ctx->slot_manager) {
                            ctx->slot_manager->stop_processing_loop();
                        }
                        ctx->disableParallelMode();
                    }
                    return [](jsi::Runtime& rt) { return jsi::Value(true); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaEnableParallelMode", enableParallelMode);

        auto queueCompletion = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaQueueCompletion"),
            4,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Object params = arguments[1].asObject(runtime);

                auto onToken = makeJsiFunction(runtime, arguments[2], callInvoker);
                auto onComplete = makeJsiFunction(runtime, arguments[3], callInvoker);

                auto ctxPtr = getContextOrThrow(contextId);
                auto originalParams = ctxPtr->params;
                parseCompletionParams(runtime, params, ctxPtr);
                common_params cparams = ctxPtr->params;
                ctxPtr->params = originalParams;

                std::vector<std::string> mediaPaths;
                if (params.hasProperty(runtime, "media_paths")) {
                    jsi::Array paths = params.getProperty(runtime, "media_paths").asObject(runtime).asArray(runtime);
                    for (size_t i = 0; i < paths.size(runtime); i++) {
                        mediaPaths.push_back(paths.getValueAtIndex(runtime, i).asString(runtime).utf8(runtime));
                    }
                }

                int chat_format = getPropertyAsInt(runtime, params, "chat_format", 0);
                std::string reasoningFormatStr = getPropertyAsString(runtime, params, "reasoning_format", "none");
                common_reasoning_format reasoning_format = common_reasoning_format_from_name(reasoningFormatStr);
                std::string generation_prompt = getPropertyAsString(runtime, params, "generation_prompt");
                std::string chat_parser = getPropertyAsString(runtime, params, "chat_parser");
                std::string prefill_text = getPropertyAsString(runtime, params, "prefill_text");
                std::string load_state_path = stripFileScheme(getPropertyAsString(runtime, params, "load_state_path"));
                std::string save_state_path = stripFileScheme(getPropertyAsString(runtime, params, "save_state_path"));
                std::string save_prompt_state_path = stripFileScheme(getPropertyAsString(runtime, params, "save_prompt_state_path"));
                int load_state_size = getPropertyAsInt(runtime, params, "load_state_size", -1);
                int save_state_size = getPropertyAsInt(runtime, params, "save_state_size", -1);

                return createPromiseTask(runtime, callInvoker, [runtimePtr = std::shared_ptr<jsi::Runtime>(&runtime, [](jsi::Runtime*){}), contextId, cparams, mediaPaths, chat_format, reasoning_format, generation_prompt, chat_parser, prefill_text, load_state_path, save_state_path, save_prompt_state_path, load_state_size, save_state_size, onToken, onComplete, callInvoker]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->parallel_mode_enabled || !ctx->slot_manager) {
                        throw std::runtime_error("Parallel mode not enabled");
                    }

                    auto tokenizeResult = ctx->tokenize(cparams.prompt, mediaPaths);
                    std::vector<llama_token> tokens = tokenizeResult.tokens;

                    auto tokenCallback = [contextId, callInvoker, ctx, runtimePtr](const rnllama::completion_token_output& token) {
                        int requestId = token.request_id;
                        rnllama::completion_chat_output parsed_output;
                        bool has_parsed_output = false;
                        if (ctx->slot_manager) {
                            auto* slot = ctx->slot_manager->get_slot_by_request_id(requestId);
                            if (slot) {
                                try {
                                    parsed_output = slot->parseChatOutput(true);
                                    has_parsed_output = true;
                                } catch (...) {
                                    has_parsed_output = false;
                                }
                            }
                        }

                        auto callbacks = RequestManager::getInstance().getRequest(contextId, requestId);
                        if (callbacks.onToken) {
                            rnllama::completion_token_output tokenCopy = token;
                            auto runtime = runtimePtr;
                            if (!runtime) {
                              return;
                            }
                            invokeAsyncTracked(callInvoker, contextId, [callbacks, contextId, tokenCopy, requestId, parsed_output, has_parsed_output, runtime](bool shouldProceed) {
                                if (!shouldProceed) return;
                                long ctxPtr = g_llamaContexts.get(contextId);
                                if (ctxPtr) {
                                    auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                                    auto& rt = *runtime;
                                    jsi::Object res = createTokenResult(rt, ctx, tokenCopy);
                                    if (has_parsed_output) {
                                        setChatOutputFields(rt, res, parsed_output);
                                    }
                                    callbacks.onToken->call(rt, res, jsi::Value(requestId));
                                }
                            });
                        }
                    };

                    auto completeCallback = [contextId, callInvoker, runtimePtr](rnllama::llama_rn_slot* slot) {
                        int requestId = slot->request_id;
                        auto callbacks = RequestManager::getInstance().takeRequest(contextId, requestId);
                        if (callbacks.onComplete) {
                            if (slot->parent_ctx && slot->ctx_sampling) {
                                common_perf_print(slot->parent_ctx->ctx, slot->ctx_sampling);
                            }

                            auto result = captureParallelCompletionResult(slot);
                            auto runtime = runtimePtr;
                            if (!runtime) {
                              return;
                            }
                            invokeAsyncTracked(callInvoker, contextId, [callbacks, contextId, result = std::move(result), runtime](bool shouldProceed) {
                                if (!shouldProceed) return;
                                long ctxPtr = g_llamaContexts.get(contextId);
                                if (!ctxPtr) {
                                    return;
                                }
                                auto ctxVal = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                                auto& rt = *runtime;
                                auto res = createParallelCompletionResult(rt, ctxVal, result);
                                callbacks.onComplete->call(rt, res);
                            });
                        }
                    };

                    int requestId = ctx->slot_manager->reserve_request_id();
                    RequestManager::getInstance().addRequest(contextId, requestId, {onToken, onComplete, nullptr});
                    try {
                        int queuedRequestId = ctx->slot_manager->queue_request(
                            cparams, tokens, mediaPaths, cparams.prompt, chat_format, reasoning_format, generation_prompt, chat_parser, prefill_text, load_state_path, save_state_path, save_prompt_state_path, load_state_size, save_state_size,
                            tokenCallback, completeCallback, requestId
                        );
                        if (queuedRequestId != requestId) {
                            RequestManager::getInstance().takeRequest(contextId, requestId);
                            throw std::runtime_error("Failed to queue completion request");
                        }
                    } catch (...) {
                        RequestManager::getInstance().takeRequest(contextId, requestId);
                        throw;
                    }

                    return [requestId](jsi::Runtime& rt) {
                        jsi::Object res(rt);
                        res.setProperty(rt, "requestId", requestId);
                        return res;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaQueueCompletion", queueCompletion);

        auto cancelRequest = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaCancelRequest"),
            2,
            [](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                int requestId = (int)arguments[1].asNumber();

                auto ctx = getContextOrThrow(contextId);
                if (ctx->slot_manager) {
                    auto result = ctx->slot_manager->cancel_request(requestId);
                    if (result == rnllama::llama_rn_cancel_result::QUEUED) {
                        auto callbacks = RequestManager::getInstance().takeRequest(contextId, requestId);
                        if (callbacks.onComplete) {
                            auto snapshot = createQueuedCancellationSnapshot(requestId);
                            auto response = createParallelCompletionResult(runtime, ctx, snapshot);
                            callbacks.onComplete->call(runtime, response);
                        }
                    }
                }

                return jsi::Value::undefined();
            }
        );
        runtime.global().setProperty(runtime, "llamaCancelRequest", cancelRequest);

        auto queueEmbedding = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaQueueEmbedding"),
            4,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string text = arguments[1].asString(runtime).utf8(runtime);
                jsi::Object params = arguments[2].asObject(runtime);
                auto onResult = makeJsiFunction(runtime, arguments[3], callInvoker);

                int embd_normalize = 0;
                bool has_embd_normalize = false;
                if (params.hasProperty(runtime, "embd_normalize")) {
                    embd_normalize = getPropertyAsInt(runtime, params, "embd_normalize", 2);
                    has_embd_normalize = true;
                }

                return createPromiseTask(runtime, callInvoker, [runtimePtr = std::shared_ptr<jsi::Runtime>(&runtime, [](jsi::Runtime*){}), contextId, text, embd_normalize, has_embd_normalize, onResult, callInvoker]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->parallel_mode_enabled || !ctx->slot_manager) {
                        throw std::runtime_error("Parallel mode not enabled");
                    }

                    const llama_vocab* vocab = llama_model_get_vocab(ctx->model);
                    const bool add_bos = llama_vocab_get_add_bos(vocab);
                    const bool is_enc_dec = llama_model_has_encoder(ctx->model);
                    std::vector<llama_token> tokens = common_tokenize(ctx->ctx, text, add_bos || is_enc_dec, true);

                    auto resultCallback = [contextId, callInvoker, runtimePtr](int32_t requestId, const std::vector<float>& embedding) {
                        auto callbacks = RequestManager::getInstance().takeRequest(contextId, requestId);
                        if (callbacks.onResult) {
                            std::vector<float> embCopy = embedding;
                            auto runtime = runtimePtr;
                            if (!runtime) {
                              return;
                            }
                            invokeAsyncTracked(callInvoker, contextId, [callbacks, embCopy, runtime](bool shouldProceed) {
                                if (!shouldProceed) return;
                                auto& rt = *runtime;
                                jsi::Array res(rt, embCopy.size());
                                for (size_t i = 0; i < embCopy.size(); i++) {
                                    res.setValueAtIndex(rt, i, (double)embCopy[i]);
                                }
                                callbacks.onResult->call(rt, res);
                            });
                        }
                    };

                    const int normalize = has_embd_normalize ? embd_normalize : ctx->params.embd_normalize;
                    int requestId = ctx->slot_manager->reserve_request_id();
                    RequestManager::getInstance().addRequest(contextId, requestId, {nullptr, nullptr, onResult});
                    try {
                        int queuedRequestId = ctx->slot_manager->queue_embedding_request(
                            tokens, normalize, resultCallback, requestId
                        );
                        if (queuedRequestId != requestId) {
                            RequestManager::getInstance().takeRequest(contextId, requestId);
                            throw std::runtime_error("Failed to queue embedding request");
                        }
                    } catch (...) {
                        RequestManager::getInstance().takeRequest(contextId, requestId);
                        throw;
                    }

                    return [requestId](jsi::Runtime& rt) {
                        jsi::Object res(rt);
                        res.setProperty(rt, "requestId", requestId);
                        return res;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaQueueEmbedding", queueEmbedding);

        auto queueRerank = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaQueueRerank"),
            5,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string query = arguments[1].asString(runtime).utf8(runtime);
                jsi::Array documentsArr = arguments[2].asObject(runtime).asArray(runtime);
                std::vector<std::string> documents;
                for (size_t i = 0; i < documentsArr.size(runtime); i++) {
                    documents.push_back(documentsArr.getValueAtIndex(runtime, i).asString(runtime).utf8(runtime));
                }
                jsi::Object params = arguments[3].asObject(runtime);
                auto onResult = makeJsiFunction(runtime, arguments[4], callInvoker);

                int normalize = getPropertyAsInt(runtime, params, "normalize", 0);

                return createPromiseTask(runtime, callInvoker, [runtimePtr = std::shared_ptr<jsi::Runtime>(&runtime, [](jsi::Runtime*){}), contextId, query, documents, normalize, onResult, callInvoker]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->parallel_mode_enabled || !ctx->slot_manager) {
                        throw std::runtime_error("Parallel mode not enabled");
                    }

                    auto resultCallback = [contextId, callInvoker, runtimePtr](int32_t requestId, const std::vector<float>& scores) {
                        auto callbacks = RequestManager::getInstance().takeRequest(contextId, requestId);
                        if (callbacks.onResult) {
                            std::vector<float> scoresCopy = scores;
                            auto runtime = runtimePtr;
                            if (!runtime) {
                              return;
                            }
                            invokeAsyncTracked(callInvoker, contextId, [callbacks, scoresCopy, runtime](bool shouldProceed) {
                                if (!shouldProceed) return;
                                auto& rt = *runtime;
                                jsi::Array res(rt, scoresCopy.size());
                                for (size_t i = 0; i < scoresCopy.size(); i++) {
                                    jsi::Object item(rt);
                                    item.setProperty(rt, "score", (double)scoresCopy[i]);
                                    item.setProperty(rt, "index", (int)i);
                                    res.setValueAtIndex(rt, i, item);
                                }
                                callbacks.onResult->call(rt, res);
                            });
                        }
                    };

                    int requestId = ctx->slot_manager->reserve_request_id();
                    RequestManager::getInstance().addRequest(contextId, requestId, {nullptr, nullptr, onResult});
                    try {
                        int queuedRequestId = ctx->slot_manager->queue_rerank_request(
                            query, documents, normalize, resultCallback, requestId
                        );
                        if (queuedRequestId != requestId) {
                            RequestManager::getInstance().takeRequest(contextId, requestId);
                            throw std::runtime_error("Failed to queue rerank request");
                        }
                    } catch (...) {
                        RequestManager::getInstance().takeRequest(contextId, requestId);
                        throw;
                    }

                    return [requestId](jsi::Runtime& rt) {
                        jsi::Object res(rt);
                        res.setProperty(rt, "requestId", requestId);
                        return res;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaQueueRerank", queueRerank);

        // Helper function to create parallel status object
        auto createParallelStatusObject = [](jsi::Runtime& rt, const rnllama::llama_rn_parallel_status& status) -> jsi::Object {
            jsi::Object result(rt);
            result.setProperty(rt, "n_parallel", status.n_parallel);
            result.setProperty(rt, "active_slots", status.active_slots);
            result.setProperty(rt, "queued_requests", status.queued_requests);

            jsi::Array requests(rt, status.requests.size());
            for (size_t i = 0; i < status.requests.size(); i++) {
                const auto& req = status.requests[i];
                jsi::Object reqObj(rt);
                reqObj.setProperty(rt, "request_id", req.request_id);
                reqObj.setProperty(rt, "type", jsi::String::createFromUtf8(rt, req.type));
                reqObj.setProperty(rt, "state", jsi::String::createFromUtf8(rt, req.state));
                reqObj.setProperty(rt, "prompt_length", (double)req.prompt_length);
                reqObj.setProperty(rt, "tokens_generated", (double)req.tokens_generated);
                reqObj.setProperty(rt, "prompt_ms", req.prompt_ms);
                reqObj.setProperty(rt, "generation_ms", req.generation_ms);
                reqObj.setProperty(rt, "tokens_per_second", req.tokens_per_second);
                requests.setValueAtIndex(rt, i, reqObj);
            }
            result.setProperty(rt, "requests", requests);

            return result;
        };

        // Get parallel status (one-time snapshot)
        auto getParallelStatus = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetParallelStatus"),
            1,
            [callInvoker, createParallelStatusObject](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();

                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->parallel_mode_enabled || !ctx->slot_manager) {
                        throw std::runtime_error("Parallel mode not enabled");
                    }

                    auto status = ctx->slot_manager->get_status();

                    return [status](jsi::Runtime& rt) {
                        jsi::Object result(rt);
                        result.setProperty(rt, "n_parallel", status.n_parallel);
                        result.setProperty(rt, "active_slots", status.active_slots);
                        result.setProperty(rt, "queued_requests", status.queued_requests);

                        jsi::Array requests(rt, status.requests.size());
                        for (size_t i = 0; i < status.requests.size(); i++) {
                            const auto& req = status.requests[i];
                            jsi::Object reqObj(rt);
                            reqObj.setProperty(rt, "request_id", req.request_id);
                            reqObj.setProperty(rt, "type", jsi::String::createFromUtf8(rt, req.type));
                            reqObj.setProperty(rt, "state", jsi::String::createFromUtf8(rt, req.state));
                            reqObj.setProperty(rt, "prompt_length", (double)req.prompt_length);
                            reqObj.setProperty(rt, "tokens_generated", (double)req.tokens_generated);
                            reqObj.setProperty(rt, "prompt_ms", req.prompt_ms);
                            reqObj.setProperty(rt, "generation_ms", req.generation_ms);
                            reqObj.setProperty(rt, "tokens_per_second", req.tokens_per_second);
                            requests.setValueAtIndex(rt, i, reqObj);
                        }
                        result.setProperty(rt, "requests", requests);

                        return result;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetParallelStatus", getParallelStatus);

        // Subscribe to parallel status changes
        auto subscribeParallelStatus = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaSubscribeParallelStatus"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                auto onStatus = makeJsiFunction(runtime, arguments[1], callInvoker);

                auto runtimePtr = std::make_shared<jsi::Runtime*>(&runtime);

                return createPromiseTask(runtime, callInvoker,
                    [contextId, onStatus, callInvoker, runtimePtr]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->parallel_mode_enabled || !ctx->slot_manager) {
                        throw std::runtime_error("Parallel mode not enabled");
                    }

                    auto statusCallback = [contextId, callInvoker, onStatus, runtimePtr](
                        const rnllama::llama_rn_parallel_status& status
                    ) {
                        // Copy status for async callback
                        rnllama::llama_rn_parallel_status statusCopy = status;

                        callInvoker->invokeAsync([onStatus, statusCopy, runtimePtr]() {
                            if (!runtimePtr || !*runtimePtr) return;
                            auto& rt = **runtimePtr;

                            jsi::Object result(rt);
                            result.setProperty(rt, "n_parallel", statusCopy.n_parallel);
                            result.setProperty(rt, "active_slots", statusCopy.active_slots);
                            result.setProperty(rt, "queued_requests", statusCopy.queued_requests);

                            jsi::Array requests(rt, statusCopy.requests.size());
                            for (size_t i = 0; i < statusCopy.requests.size(); i++) {
                                const auto& req = statusCopy.requests[i];
                                jsi::Object reqObj(rt);
                                reqObj.setProperty(rt, "request_id", req.request_id);
                                reqObj.setProperty(rt, "type", jsi::String::createFromUtf8(rt, req.type));
                                reqObj.setProperty(rt, "state", jsi::String::createFromUtf8(rt, req.state));
                                reqObj.setProperty(rt, "prompt_length", (double)req.prompt_length);
                                reqObj.setProperty(rt, "tokens_generated", (double)req.tokens_generated);
                                reqObj.setProperty(rt, "prompt_ms", req.prompt_ms);
                                reqObj.setProperty(rt, "generation_ms", req.generation_ms);
                                reqObj.setProperty(rt, "tokens_per_second", req.tokens_per_second);
                                requests.setValueAtIndex(rt, i, reqObj);
                            }
                            result.setProperty(rt, "requests", requests);

                            onStatus->call(rt, result);
                        });
                    };

                    int32_t subscriberId = ctx->slot_manager->add_status_subscriber(statusCallback);

                    return [subscriberId](jsi::Runtime& rt) {
                        jsi::Object res(rt);
                        res.setProperty(rt, "subscriberId", subscriberId);
                        return res;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaSubscribeParallelStatus", subscribeParallelStatus);

        // Unsubscribe from parallel status changes
        auto unsubscribeParallelStatus = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaUnsubscribeParallelStatus"),
            2,
            [](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                int subscriberId = (int)arguments[1].asNumber();

                long ctxPtr = g_llamaContexts.get(contextId);
                if (ctxPtr) {
                    auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                    if (ctx->slot_manager) {
                        ctx->slot_manager->remove_status_subscriber(subscriberId);
                    }
                }

                return jsi::Value::undefined();
            }
        );
        runtime.global().setProperty(runtime, "llamaUnsubscribeParallelStatus", unsubscribeParallelStatus);

        auto releaseContext = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaReleaseContext"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                 int contextId = (int)arguments[0].asNumber();
                 return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                     RequestManager::getInstance().clearContext(contextId);
                     long ctxPtr = g_llamaContexts.get(contextId);
                     if (ctxPtr) {
                         auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                         if (ctx->completion) {
                             ctx->completion->is_interrupted = true;
                         }
                         if (ctx->slot_manager) {
                             ctx->slot_manager->stop_processing_loop();
                         }
                     }

                     // Wait for ALL other tasks on this context to complete (including their
                     // invokeAsync callbacks) before deleting. This prevents race conditions
                     // where we delete the context while a completion's JS callback is still
                     // accessing ctx->completion.
                     TaskManager::getInstance().waitForContext(contextId, 0);
                     if (TaskManager::getInstance().isShuttingDown()) {
                         return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                     }

                     if (ctxPtr) {
                         auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                         // Remove from map FIRST, then delete.
                         // This ensures any concurrent lookups via g_llamaContexts.get()
                         // will return 0 (not found) rather than a dangling pointer.
                         removeContext(contextId);
                         delete ctx;
                     }
                     return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                 }, contextId, false);  // trackTask=false - release should not count itself
            }
        );
        runtime.global().setProperty(runtime, "llamaReleaseContext", releaseContext);

        auto releaseAllContexts = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaReleaseAllContexts"),
            0,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                 return createPromiseTask(runtime, callInvoker, []() -> PromiseResultGenerator {
                     RequestManager::getInstance().clearAll();

                     auto contexts = g_llamaContexts.snapshot();
                     for (const auto& entry : contexts) {
                         long ctxPtr = entry.second;
                         if (!ctxPtr) {
                             continue;
                         }
                         auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
                         if (ctx->completion) {
                             ctx->completion->is_interrupted = true;
                         }
                         if (ctx->slot_manager) {
                             ctx->slot_manager->stop_processing_loop();
                         }
                     }

                     // Wait for ALL tasks to complete (including their invokeAsync callbacks)
                     TaskManager::getInstance().waitForAll(0);
                     if (TaskManager::getInstance().isShuttingDown()) {
                         return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                     }

                     g_llamaContexts.clear([](long ptr) {
                        if (ptr) {
                            auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ptr);
                            delete ctx;
                        }
                     });
                     return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                 }, -1, false);  // contextId=-1 (not tracked), trackTask=false
            }
        );
        runtime.global().setProperty(runtime, "llamaReleaseAllContexts", releaseAllContexts);

        auto setContextLimitFn = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaSetContextLimit"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int64_t limit = (int64_t)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [limit]() -> PromiseResultGenerator {
                    setContextLimit(limit);
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                });
            }
        );
        runtime.global().setProperty(runtime, "llamaSetContextLimit", setContextLimitFn);

        // LoRA Adapters
        auto applyLoraAdapters = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaApplyLoraAdapters"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Array loraList = arguments[1].asObject(runtime).asArray(runtime);
                std::vector<common_adapter_lora_info> lora_adapters;
                for (size_t i = 0; i < loraList.size(runtime); i++) {
                    jsi::Object item = loraList.getValueAtIndex(runtime, i).asObject(runtime);
                    common_adapter_lora_info la;
                    la.path = getPropertyAsString(runtime, item, "path");
                    la.scale = getPropertyAsFloat(runtime, item, "scaled", 1.0f);
                    if (!la.path.empty()) {
                        lora_adapters.push_back(la);
                    }
                }

                return createPromiseTask(runtime, callInvoker, [contextId, lora_adapters]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    ctx->applyLoraAdapters(lora_adapters);
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaApplyLoraAdapters", applyLoraAdapters);

        auto removeLoraAdapters = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaRemoveLoraAdapters"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    ctx->removeLoraAdapters();
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaRemoveLoraAdapters", removeLoraAdapters);

        auto getLoadedLoraAdapters = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetLoadedLoraAdapters"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    auto adapters = ctx->getLoadedLoraAdapters();
                    return [adapters](jsi::Runtime& rt) {
                        jsi::Array res(rt, adapters.size());
                        for (size_t i = 0; i < adapters.size(); i++) {
                            jsi::Object item(rt);
                            item.setProperty(rt, "path", jsi::String::createFromUtf8(rt, adapters[i].path));
                            item.setProperty(rt, "scaled", (double)adapters[i].scale);
                            res.setValueAtIndex(rt, i, item);
                        }
                        return res;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetLoadedLoraAdapters", getLoadedLoraAdapters);

        // Multimodal
        auto initMultimodal = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaInitMultimodal"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Object params = arguments[1].asObject(runtime);
                std::string path = getPropertyAsString(runtime, params, "path");
                bool use_gpu = getPropertyAsBool(runtime, params, "use_gpu", true);
                int image_min_tokens = getPropertyAsInt(runtime, params, "image_min_tokens", -1);
                int image_max_tokens = getPropertyAsInt(runtime, params, "image_max_tokens", -1);

                return createPromiseTask(runtime, callInvoker, [contextId, path, use_gpu, image_min_tokens, image_max_tokens]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    bool result = ctx->initMultimodal(path, use_gpu, image_min_tokens, image_max_tokens);
                    return [result](jsi::Runtime& rt) { return jsi::Value(result); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaInitMultimodal", initMultimodal);

        auto isMultimodalEnabled = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaIsMultimodalEnabled"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    bool result = ctx->isMultimodalEnabled();
                    return [result](jsi::Runtime& rt) { return jsi::Value(result); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaIsMultimodalEnabled", isMultimodalEnabled);

        auto getMultimodalSupport = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetMultimodalSupport"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isMultimodalEnabled()) throw std::runtime_error("Multimodal is not enabled");
                    bool vision = ctx->isMultimodalSupportVision();
                    bool audio = ctx->isMultimodalSupportAudio();
                    return [vision, audio](jsi::Runtime& rt) {
                        jsi::Object res(rt);
                            res.setProperty(rt, "vision", vision);
                            res.setProperty(rt, "audio", audio);
                            return res;
                        };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetMultimodalSupport", getMultimodalSupport);

        auto releaseMultimodal = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaReleaseMultimodal"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    ctx->releaseMultimodal();
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaReleaseMultimodal", releaseMultimodal);

        // Vocoder
        auto initVocoder = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaInitVocoder"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Object params = arguments[1].asObject(runtime);
                std::string path = getPropertyAsString(runtime, params, "path");
                int n_batch = getPropertyAsInt(runtime, params, "n_batch", 512);
                // use_gpu defaults to follow the main context's n_gpu_layers
                // (any > 0 means the backbone is GPU-offloaded — pair the
                // codec / codec_lm there too unless the caller overrides).
                bool use_gpu_default = false;
                {
                    auto ctx_for_default = getContextOrThrow(contextId);
                    use_gpu_default = ctx_for_default->params.n_gpu_layers > 0;
                }
                bool use_gpu = getPropertyAsBool(runtime, params, "use_gpu", use_gpu_default);

                return createPromiseTask(runtime, callInvoker, [contextId, path, n_batch, use_gpu]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    bool result = ctx->initVocoder(path, n_batch, use_gpu);
                    return [result](jsi::Runtime& rt) { return jsi::Value(result); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaInitVocoder", initVocoder);

        auto isVocoderEnabled = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaIsVocoderEnabled"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    bool result = ctx->isVocoderEnabled();
                    return [result](jsi::Runtime& rt) { return jsi::Value(result); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaIsVocoderEnabled", isVocoderEnabled);

        auto getFormattedAudioCompletion = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetFormattedAudioCompletion"),
            4,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string speakerJsonStr = arguments[1].asString(runtime).utf8(runtime);
                std::string textToSpeak = arguments[2].asString(runtime).utf8(runtime);
                // Optional 4th arg: speakerId (registry id >= 0, or -1 for none).
                int speakerId = (count >= 4 && arguments[3].isNumber())
                    ? (int)arguments[3].asNumber()
                    : -1;

                return createPromiseTask(runtime, callInvoker, [contextId, speakerJsonStr, textToSpeak, speakerId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    try {
                        auto audio_result = ctx->tts_wrapper->getFormattedAudioCompletion(ctx, speakerJsonStr, textToSpeak, speakerId);
                        return [audio_result](jsi::Runtime& rt) {
                            jsi::Object res(rt);
                            res.setProperty(rt, "prompt", jsi::String::createFromUtf8(rt, audio_result.prompt));
                            if (!audio_result.grammar.empty()) {
                                res.setProperty(rt, "grammar", jsi::String::createFromUtf8(rt, audio_result.grammar));
                            }
                            res.setProperty(rt, "embedding", audio_result.embedding);
                            res.setProperty(rt, "flow", jsi::String::createFromUtf8(rt, audio_result.flow));
                            return res;
                        };
                    } catch (const std::exception &e) {
                        throw std::runtime_error(e.what());
                    }
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetFormattedAudioCompletion", getFormattedAudioCompletion);

        auto getTTSCapabilities = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetTTSCapabilities"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");
                    auto cap = ctx->tts_wrapper->getTTSCapabilities(ctx);
                    return [cap](jsi::Runtime& rt) {
                        jsi::Object obj(rt);
                        obj.setProperty(rt, "type", jsi::Value(cap.type));
                        obj.setProperty(rt, "promptKind", jsi::String::createFromUtf8(rt, cap.prompt_kind));
                        obj.setProperty(rt, "family", jsi::String::createFromUtf8(rt, cap.family));
                        obj.setProperty(rt, "requiresPhonemes", jsi::Value(cap.requires_phonemes));
                        obj.setProperty(rt, "defaultLanguage", jsi::String::createFromUtf8(rt, cap.default_language));
                        return obj;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetTTSCapabilities", getTTSCapabilities);

        auto decodeAudioTokens = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaDecodeAudioTokens"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Array tokensArr = arguments[1].asObject(runtime).asArray(runtime);
                std::vector<llama_token> tokens;
                for (size_t i = 0; i < tokensArr.size(runtime); i++) {
                    tokens.push_back((llama_token)tokensArr.getValueAtIndex(runtime, i).asNumber());
                }

                return createPromiseTask(runtime, callInvoker, [contextId, tokens]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    try {
                        auto audio_data = ctx->tts_wrapper->decodeAudioTokens(ctx, tokens);
                        return [audio_data](jsi::Runtime& rt) {
                            jsi::Array res(rt, audio_data.size());
                            for (size_t i = 0; i < audio_data.size(); i++) {
                                res.setValueAtIndex(rt, i, (double)audio_data[i]);
                            }
                            return res;
                        };
                    } catch (const std::exception &e) {
                        throw std::runtime_error(e.what());
                    }
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaDecodeAudioTokens", decodeAudioTokens);

        // generateAudioCodes — drives the backbone + codec_lm AR loop for
        // codec_lm-flow models (CSM, etc.).  Args:
        //   (contextId, optsJson, onFrame?)
        // optsJson: { prompt, maxFrames?, temperature?, topP?, topK?, seed? }
        // onFrame:  optional (step:number, codes:number[]) => void — fired
        //           per-frame as audio codes are produced.
        // Returns { codes:number[], nCodebook, nFrames, stoppedOnEos, aborted }.
        auto generateAudioCodes = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGenerateAudioCodes"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string optsJson = arguments[1].asString(runtime).utf8(runtime);

                std::shared_ptr<jsi::Function> onFrame;
                if (count >= 3 && arguments[2].isObject() &&
                    arguments[2].asObject(runtime).isFunction(runtime)) {
                    onFrame = std::make_shared<jsi::Function>(
                        arguments[2].asObject(runtime).asFunction(runtime));
                }
                jsi::Runtime * runtimePtr = &runtime;

                return createPromiseTask(runtime, callInvoker, [contextId, optsJson, onFrame, runtimePtr, callInvoker]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    rnllama::llama_rn_audio_codes_options opts;
                    try {
                        auto j = nlohmann::ordered_json::parse(optsJson);
                        opts.prompt      = j.value("prompt", std::string());
                        opts.max_frames  = j.value("maxFrames",   500);
                        opts.temperature = j.value("temperature", 0.9f);
                        opts.top_p       = j.value("topP",        0.95f);
                        opts.top_k       = j.value("topK",        50);
                        opts.seed        = j.value("seed",        0u);

                    } catch (const std::exception &e) {
                        throw std::runtime_error(std::string("invalid options JSON: ") + e.what());
                    }
                    if (opts.prompt.empty()) {
                        throw std::runtime_error("generateAudioCodes: prompt is empty");
                    }

                    rnllama::llama_rn_audio_codes_progress_cb cb;
                    if (onFrame) {
                        // Fire-and-forget per-frame notification.  We never
                        // block on the JS side, so the return value is
                        // always "continue"; aborting from JS isn't wired
                        // through here yet.
                        cb = [onFrame, runtimePtr, callInvoker](int step, const std::vector<int32_t> &codes) -> bool {
                            std::vector<int32_t> codes_copy = codes;
                            callInvoker->invokeAsync([onFrame, runtimePtr, step, codes_copy]() {
                                auto &rt = *runtimePtr;
                                jsi::Array arr(rt, codes_copy.size());
                                for (size_t i = 0; i < codes_copy.size(); ++i) {
                                    arr.setValueAtIndex(rt, i, (double) codes_copy[i]);
                                }
                                onFrame->call(rt, jsi::Value((double) step), arr);
                            });
                            return true;
                        };
                    }

                    try {
                        auto r = ctx->tts_wrapper->generateAudioCodes(ctx, opts, cb);
                        return [r](jsi::Runtime& rt) {
                            jsi::Object obj(rt);
                            jsi::Array arr(rt, r.codes.size());
                            for (size_t i = 0; i < r.codes.size(); ++i) {
                                arr.setValueAtIndex(rt, i, (double) r.codes[i]);
                            }
                            obj.setProperty(rt, "codes", arr);
                            obj.setProperty(rt, "nCodebook",     jsi::Value((double) r.n_codebook));
                            obj.setProperty(rt, "nFrames",       jsi::Value((double) r.n_frames));
                            obj.setProperty(rt, "stoppedOnEos",  jsi::Value(r.stopped_on_eos));
                            obj.setProperty(rt, "aborted",       jsi::Value(r.aborted));
                            return obj;
                        };
                    } catch (const std::exception &e) {
                        throw std::runtime_error(e.what());
                    }
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGenerateAudioCodes", generateAudioCodes);

        // llamaCreateSpeaker(ctxId, optsJson)
        //   optsJson: { pcm: number[], inputSampleRate: number, refText: string,
        //               bake: boolean, emotion?: number }
        // Resolves: { id: number, family: string, rows: number, baked: boolean }
        auto createSpeaker = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaCreateSpeaker"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                std::string optsJson = arguments[1].asString(runtime).utf8(runtime);

                std::vector<float> pcm;
                int inputSampleRate = 0;
                std::string refText;
                float emotion = 0.5f;
                bool has_emotion = false;
                bool bake = false;

                try {
                    auto j = nlohmann::ordered_json::parse(optsJson);
                    if (j.contains("pcm") && j["pcm"].is_array()) {
                        pcm.reserve(j["pcm"].size());
                        for (const auto & v : j["pcm"]) {
                            pcm.push_back((float) v.get<double>());
                        }
                    }
                    inputSampleRate = j.value("inputSampleRate", 0);
                    refText         = j.value("refText", std::string());
                    bake            = j.value("bake", false);
                    if (j.contains("emotion") && j["emotion"].is_number()) {
                        has_emotion = true;
                        emotion = (float) j["emotion"].get<double>();
                    }
                } catch (const std::exception & e) {
                    throw std::runtime_error(std::string("createSpeaker: invalid options JSON: ") + e.what());
                }

                return createPromiseTask(runtime, callInvoker, [contextId, pcm, inputSampleRate, refText, emotion, has_emotion, bake]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    auto cap = ctx->tts_wrapper->getTTSCapabilities(ctx);
                    std::string family = cap.family;

                    const int speakerId = ctx->tts_wrapper->createSpeaker(
                        ctx, pcm, inputSampleRate, refText, emotion, has_emotion, bake);

                    const rnllama::rn_speaker * spk = ctx->tts_wrapper->getSpeaker(speakerId);
                    int rows  = spk ? spk->rows  : 0;
                    bool baked = spk ? spk->baked : false;

                    return [speakerId, family, rows, baked](jsi::Runtime& rt) {
                        jsi::Object obj(rt);
                        obj.setProperty(rt, "id",     jsi::Value((double) speakerId));
                        obj.setProperty(rt, "family", jsi::String::createFromUtf8(rt, family));
                        obj.setProperty(rt, "rows",   jsi::Value((double) rows));
                        obj.setProperty(rt, "baked",  jsi::Value(baked));
                        return obj;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaCreateSpeaker", createSpeaker);

        // llamaBakeSpeaker(ctxId, speakerId)
        // Resolves: { rows: number, baked: boolean }
        auto bakeSpeaker = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaBakeSpeaker"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                int speakerId = (int)arguments[1].asNumber();

                return createPromiseTask(runtime, callInvoker, [contextId, speakerId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    ctx->tts_wrapper->bakeSpeaker(ctx, speakerId);

                    const rnllama::rn_speaker * spk = ctx->tts_wrapper->getSpeaker(speakerId);
                    if (!spk) throw std::runtime_error("bakeSpeaker: speaker id not found");

                    int rows  = spk->rows;
                    bool baked = spk->baked;

                    return [rows, baked](jsi::Runtime& rt) {
                        jsi::Object obj(rt);
                        obj.setProperty(rt, "rows",  jsi::Value((double) rows));
                        obj.setProperty(rt, "baked", jsi::Value(baked));
                        return obj;
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaBakeSpeaker", bakeSpeaker);

        // llamaReleaseSpeaker(ctxId, speakerId)
        // Resolves: void
        auto releaseSpeaker = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaReleaseSpeaker"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                int speakerId = (int)arguments[1].asNumber();

                return createPromiseTask(runtime, callInvoker, [contextId, speakerId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    ctx->tts_wrapper->releaseSpeaker(speakerId);

                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaReleaseSpeaker", releaseSpeaker);

        auto decodeAudioEmbeddings = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaDecodeAudioEmbeddings"),
            3,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                jsi::Array embeddingsArr = arguments[1].asObject(runtime).asArray(runtime);
                int embeddingDim = (int)arguments[2].asNumber();
                std::vector<float> embeddings;
                embeddings.reserve(embeddingsArr.size(runtime));
                for (size_t i = 0; i < embeddingsArr.size(runtime); i++) {
                    embeddings.push_back((float)embeddingsArr.getValueAtIndex(runtime, i).asNumber());
                }

                return createPromiseTask(runtime, callInvoker, [contextId, embeddings, embeddingDim]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    try {
                        auto audio_data = ctx->tts_wrapper->decodeAudioEmbeddings(ctx, embeddings, embeddingDim);
                        return [audio_data](jsi::Runtime& rt) {
                            jsi::Array res(rt, audio_data.size());
                            for (size_t i = 0; i < audio_data.size(); i++) {
                                res.setValueAtIndex(rt, i, (double)audio_data[i]);
                            }
                            return res;
                        };
                    } catch (const std::exception &e) {
                        throw std::runtime_error(e.what());
                    }
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaDecodeAudioEmbeddings", decodeAudioEmbeddings);

        auto getAudioSampleRate = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaGetAudioSampleRate"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();

                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    if (!ctx->isVocoderEnabled()) throw std::runtime_error("Vocoder is not enabled");

                    const int sample_rate = ctx->tts_wrapper->getAudioSampleRate();
                    return [sample_rate](jsi::Runtime& rt) {
                        return jsi::Value((double)sample_rate);
                    };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaGetAudioSampleRate", getAudioSampleRate);

        // Cache management
        auto clearCache = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaClearCache"),
            2,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                bool clearData = count > 1 && arguments[1].isBool() ? arguments[1].asBool() : false;
                return createPromiseTask(runtime, callInvoker, [contextId, clearData]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    ctx->clearCache(clearData);
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaClearCache", clearCache);

        auto releaseVocoder = jsi::Function::createFromHostFunction(runtime,
            jsi::PropNameID::forAscii(runtime, "llamaReleaseVocoder"),
            1,
            [callInvoker](jsi::Runtime& runtime, const jsi::Value& thisValue, const jsi::Value* arguments, size_t count) -> jsi::Value {
                int contextId = (int)arguments[0].asNumber();
                return createPromiseTask(runtime, callInvoker, [contextId]() -> PromiseResultGenerator {
                    auto ctx = getContextOrThrow(contextId);
                    throwIfContextBusy(ctx);
                    ctx->releaseVocoder();
                    return [](jsi::Runtime& rt) { return jsi::Value::undefined(); };
                }, contextId);
            }
        );
        runtime.global().setProperty(runtime, "llamaReleaseVocoder", releaseVocoder);
    }

    void cleanupJSIBindings() {
        TaskManager::getInstance().beginShutdown();
        {
            std::lock_guard<std::mutex> lock(g_log_mutex);
            g_log_handler.reset();
            g_log_invoker.reset();
            g_log_runtime.reset();
        }
        llama_log_set(llama_log_callback_default, nullptr);

        RequestManager::getInstance().clearAll();
        auto contexts = g_llamaContexts.snapshot();
        for (const auto& entry : contexts) {
            long ctxPtr = entry.second;
            if (!ctxPtr) {
                continue;
            }
            auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ctxPtr);
            if (ctx->completion) {
                ctx->completion->is_interrupted = true;
            }
            if (ctx->slot_manager) {
                ctx->slot_manager->stop_processing_loop();
            }
        }

        if (contexts.empty()) {
            g_context_limit.store(-1);
            return;
        }
        ThreadPool::getInstance().shutdown();

        g_llamaContexts.clear([](long ptr) {
            if (ptr) {
                auto ctx = reinterpret_cast<rnllama::llama_rn_context*>(ptr);
                delete ctx;
            }
        });
        g_context_limit.store(-1);
    }
}