/**
 * T18 端侧小模型 SPIKE · N-API 桥接
 *
 * ArkTS 侧用法：import localllm from 'liblocalllm.so'
 * 耗时接口（loadModel/generate）一律 async work + Promise，不阻塞 UI 线程。
 */
#include "local_llm.h"

#include <napi/native_api.h>
#include <string>

namespace {

struct AsyncTask {
    napi_async_work work = nullptr;
    napi_deferred deferred = nullptr;
    std::string arg;       // loadModel: modelDir；generate: prompt
    int maxNewTokens = 256; // 仅 generate 用
    std::string result;    // 返回值（空串或 JSON）
    bool isLoad = true;    // true=loadModel，false=generate
};

napi_value ResolveString(napi_env env, napi_deferred deferred, const std::string &value) {
    napi_value js;
    napi_create_string_utf8(env, value.c_str(), value.size(), &js);
    napi_resolve_deferred(env, deferred, js);
    return js;
}

void TaskExecute(napi_env /*env*/, void *data) {
    auto *task = static_cast<AsyncTask *>(data);
    if (task->isLoad) {
        task->result = LocalLlm::instance().load(task->arg); // 空串=成功
    } else {
        task->result = LocalLlm::instance().generate(task->arg, task->maxNewTokens);
    }
}

void TaskComplete(napi_env env, napi_status /*status*/, void *data) {
    auto *task = static_cast<AsyncTask *>(data);
    ResolveString(env, task->deferred, task->result);
    napi_delete_async_work(env, task->work);
    delete task;
}

napi_value RunAsync(napi_env env, napi_callback_info info, bool isLoad) {
    size_t argc = 2;
    napi_value argv[2] = {nullptr, nullptr};
    napi_get_cb_info(env, info, &argc, argv, nullptr, nullptr);

    auto *task = new AsyncTask();
    task->isLoad = isLoad;

    char buf[8192];
    size_t len = 0;
    napi_get_value_string_utf8(env, argv[0], buf, sizeof(buf), &len);
    task->arg.assign(buf, len);
    if (!isLoad && argc >= 2) {
        napi_get_value_int32(env, argv[1], &task->maxNewTokens);
    }

    napi_value promise;
    napi_create_promise(env, &task->deferred, &promise);
    napi_value name;
    napi_create_string_utf8(env, isLoad ? "loadModel" : "generate", NAPI_AUTO_LENGTH, &name);
    napi_create_async_work(env, nullptr, name, TaskExecute, TaskComplete, task, &task->work);
    napi_queue_async_work(env, task->work);
    return promise;
}

napi_value LoadModel(napi_env env, napi_callback_info info) {
    return RunAsync(env, info, true);
}

napi_value Generate(napi_env env, napi_callback_info info) {
    return RunAsync(env, info, false);
}

napi_value UnloadModel(napi_env env, napi_callback_info /*info*/) {
    LocalLlm::instance().unload();
    napi_value ret;
    napi_get_undefined(env, &ret);
    return ret;
}

napi_value IsLoaded(napi_env env, napi_callback_info /*info*/) {
    napi_value ret;
    napi_get_boolean(env, LocalLlm::instance().loaded(), &ret);
    return ret;
}

napi_value Init(napi_env env, napi_value exports) {
    napi_property_descriptor desc[] = {
        {"loadModel", nullptr, LoadModel, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"generate", nullptr, Generate, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"unloadModel", nullptr, UnloadModel, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"isLoaded", nullptr, IsLoaded, nullptr, nullptr, nullptr, napi_default, nullptr},
    };
    napi_define_properties(env, exports, sizeof(desc) / sizeof(desc[0]), desc);
    return exports;
}

} // namespace

static napi_module localllmModule = {
    .nm_version = 1,
    .nm_flags = 0,
    .nm_filename = nullptr,
    .nm_register_func = Init,
    .nm_modname = "localllm",
    .nm_priv = nullptr,
    .reserved = {0},
};

extern "C" __attribute__((constructor)) void RegisterLocalllmModule() {
    napi_module_register(&localllmModule);
}