/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "coretextembeddingservice.h"

#include <jsoncpp/json/json.h>
#include <sodium.h>

#include "utils/logger.h"

const std::string CoreTextEmbeddingService::objectPath_ = "/com/kylin/AiRuntime/CoreTextEmbeddingService";

std::set<int> CoreTextEmbeddingService::sessionIdSet;

namespace {
std::string createResponse(const std::vector<float> &result, int errorCode, const std::string &errorMessage) {
    Json::Value root;
    root["errorCode"] = errorCode;
    root["errorMessage"] = errorMessage;

    Json::Value resultArray;
    for (int i = 0; i < result.size(); ++i) {
        resultArray.append(result[i]);
    }
    root["vector_result"] = resultArray;
    return root.toStyledString();
}

}  // namespace

CoreTextEmbeddingService::CoreTextEmbeddingService(GDBusConnection &connection, EngineSelector &engineSelector)
    : connection_(connection), engineSelector_(engineSelector) {
    exportSkeleton();
}

CoreTextEmbeddingService::~CoreTextEmbeddingService() {
    unexportSkeleton();
    LOGD("~CoreTextEmbeddingService");
}

void CoreTextEmbeddingService::stop() {
    for (const auto &handlerId : handlerIds_) {
        g_signal_handler_disconnect(delegate_, handlerId);
    }
    handlerIds_.clear();
}

void CoreTextEmbeddingService::exportSkeleton() {
    delegate_ = ai_runtime_core_text_embedding_service_skeleton_new();

    gulong handlerId{};
    handlerId = g_signal_connect(delegate_, "handle-init", G_CALLBACK(onHandleInitEngine), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle-embedding-text", G_CALLBACK(onHandleEmbeddingText), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle-get-model-info", G_CALLBACK(onHandleGetModelInfo), this);
    handlerIds_.push_back(handlerId);

    GError *error = nullptr;

    isExported_ = g_dbus_interface_skeleton_export(G_DBUS_INTERFACE_SKELETON(delegate_), &connection_,
                                                   objectPath_.c_str(), &error);
    if (!isExported_) {
        LOGE("Error creating server at address {} {}", objectPath_, error->message);
        g_error_free(error);
    }
}

void CoreTextEmbeddingService::unexportSkeleton() {
    if (delegate_ == nullptr) {
        return;
    }

    if (isExported_) {
        g_dbus_interface_skeleton_unexport_from_connection(G_DBUS_INTERFACE_SKELETON(delegate_), &connection_);
    }
    sessionIdSet.clear();
    g_object_unref(delegate_);
}

bool CoreTextEmbeddingService::onHandleInitEngine(AiRuntimeCoreTextEmbeddingService *delegate,
                                                  GDBusMethodInvocation *invocation, const gchar *inputJson,
                                                  gpointer userData) {
    std::shared_ptr<CoreTextEmbeddingService> textEmbeddingService =
        static_cast<CoreTextEmbeddingService *>(userData)->getSharedObject();

    if (!textEmbeddingService) {
        LOGE("[onHandleInitEngine] error: textEmbeddingService is nullptr!");
        ai_runtime_core_text_embedding_service_complete_init(delegate, invocation, -1, -1);
        return true;
    }

    std::string engineName;
    // parse input json
    Json::Value root;
    Json::Reader r;
    bool isReader = r.parse(inputJson, root);

    if (!isReader) {
        LOGE("[onHandleInitEngine] error: parse input json failed!");
        ai_runtime_core_text_embedding_service_complete_init(delegate, invocation, -1, -1);
        return true;
    }
    std::string key = "engineName";
    if (root.isMember(key)) {
        engineName = root[key].asString();
    } else {
        LOGE("[onHandleInitEngine] error: Key: {} does not exist in the JSON object!", key);
        ai_runtime_core_text_embedding_service_complete_init(delegate, invocation, -1, -1);
        return true;
    }

    int sessionId = textEmbeddingService->generateRandomSessionId();

    ModelConfigManager::getInstance().waitForModelInfosLoaded();
    bool initResult = textEmbeddingService->initEngine(engineName, sessionId);

    if (!initResult) {
        ai_runtime_core_text_embedding_service_complete_init(delegate, invocation, -1, -1);
        return true;
    }
    LOGI("[onHandleInitEngine] service: The sessionId of the engine is {}", sessionId);
    LOGI("[onHandleInitEngine] start load text embedding model.");
    ai_engine::result::AiEngineBaseResult res =
        textEmbeddingService->textEmbeddingEngineMap_[sessionId]->initTextEmbeddingModule();
    if (!res.success) {
        LOGE("[onHandleInitEngine] load text embedding model error: {}", res.error.getMessage());
        ai_runtime_core_text_embedding_service_complete_init(delegate, invocation, sessionId, -1);
        return true;
    }
    LOGI("[onHandleInitEngine] load text embedding model success.");
    int initModuleErrorCode = 0;
    ai_runtime_core_text_embedding_service_complete_init(delegate, invocation, sessionId, initModuleErrorCode);

    return true;
}

bool CoreTextEmbeddingService::initEngine(const std::string &engineName, gint sessionId) {
    auto textEmbeddingEngine = engineSelector_.selectCoreAiEmbeddingEngine(engineName);
    if (!textEmbeddingEngine) {
        LOGE("[initEngine]: textEmbeddingEngine is nullptr! engineName: {}", engineName);
        return false;
    }
    textEmbeddingEngineMap_[sessionId] = std::move(textEmbeddingEngine);

    return true;
}

int CoreTextEmbeddingService::generateRandomSessionId() {
    uint32_t sessionId;
    randombytes_buf(&sessionId, sizeof(sessionId));
    sessionId = 1 + (sessionId & INT_MAX);
    while (sessionIdSet.find(sessionId) != sessionIdSet.end()) {
        randombytes_buf(&sessionId, sizeof(sessionId));
        sessionId = 1 + (sessionId & INT_MAX);
    }
    sessionIdSet.insert(sessionId);
    return sessionId;
}

bool CoreTextEmbeddingService::onHandleEmbeddingText(AiRuntimeCoreTextEmbeddingService *delegate,
                                                     GDBusMethodInvocation *invocation, const gchar *inputJson,
                                                     gpointer userData) {
    std::shared_ptr<CoreTextEmbeddingService> textEmbeddingService =
        static_cast<CoreTextEmbeddingService *>(userData)->getSharedObject();
    if (!textEmbeddingService) {
        LOGE("[onHandleEmbeddingText] error: textEmbeddingService is nullptr!");

        std::string response =
            createResponse({}, -1, "[onHandleEmbeddingText] error: textEmbeddingService is nullptr!");
        ai_runtime_core_text_embedding_service_complete_embedding_text(delegate, invocation, response.c_str());
        return true;
    }

    // parse input json
    std::string input(inputJson);
    std::string text;
    int sessionId;
    if (!textEmbeddingService->parseTextEmbeddingInput(input, text, sessionId)) {
        LOGE("[onHandleEmbeddingText] error: parse input json failed!");
        std::string resultJson = createResponse({}, -1, "[onHandleEmbeddingText] error: parse input json failed!");
        ai_runtime_core_text_embedding_service_complete_embedding_text(delegate, invocation, resultJson.c_str());
        return true;
    }

    if (textEmbeddingService->textEmbeddingEngineMap_.find(sessionId) ==
        textEmbeddingService->textEmbeddingEngineMap_.end()) {
        LOGE("[onHandleEmbeddingText] error: TextEmbeddingEngine is nullptr!");
        std::string resultJson =
            createResponse({}, -1, "[onHandleEmbeddingText] error: TextEmbeddingEngine is nullptr!");
        ai_runtime_core_text_embedding_service_complete_embedding_text(delegate, invocation, resultJson.c_str());
        return true;
    }

    auto callback = [delegate, invocation](const ai_engine::core_ai::embedding::EmbeddingResult &result) {
        {
            if (result.embeddings.empty()) {
                std::string errorMessage = result.error.getMessage();
                LOGE("[onHandleEmbeddingText] embedding engine error: errorMessage : {}", errorMessage);
                std::string resultJson = createResponse({}, -1,
                                                        "[onHandleEmbeddingText] embedding engine error: "
                                                        "errorMessage : " +
                                                            errorMessage);
                ai_runtime_core_text_embedding_service_complete_embedding_text(delegate, invocation,
                                                                               resultJson.c_str());
                return;
            }
            std::string resultJson = createResponse(result.embeddings, 0, "Success");

            ai_runtime_core_text_embedding_service_complete_embedding_text(delegate, invocation, resultJson.c_str());
        }
    };

    // 目前is_question模型暂未支持
    textEmbeddingService->textEmbeddingEngineMap_[sessionId]->setEmbeddingResultCallback(callback);
    textEmbeddingService->textEmbeddingEngineMap_[sessionId]->embeddingText(text);
    return true;
}

bool CoreTextEmbeddingService::onHandleGetModelInfo(AiRuntimeCoreTextEmbeddingService *delegate,
                                                    GDBusMethodInvocation *invocation, int sessionId,
                                                    gpointer userData) {
    std::shared_ptr<CoreTextEmbeddingService> textEmbeddingService =
        static_cast<CoreTextEmbeddingService *>(userData)->getSharedObject();

    if (!textEmbeddingService) {
        LOGE("[onHandleGetModelInfo] error: textEmbeddingService is nullptr!");
        std::string res = "";
        ai_runtime_core_text_embedding_service_complete_get_model_info(delegate, invocation, res.c_str());
        return true;
    }

    if (textEmbeddingService->textEmbeddingEngineMap_.find(sessionId) ==
        textEmbeddingService->textEmbeddingEngineMap_.end()) {
        LOGE("[onHandleGetModelInfo] error: TextEmbeddingEngine is nullptr!");
        std::string res = "";
        ai_runtime_core_text_embedding_service_complete_get_model_info(delegate, invocation, res.c_str());
        return true;
    }
    std::string modelInfo = textEmbeddingService->textEmbeddingEngineMap_[sessionId]->modelInfo();
    if (modelInfo.empty()) {
        LOGE("[onHandleGetModelInfo] error: modelInfo is empty!");
        ai_runtime_core_text_embedding_service_complete_get_model_info(delegate, invocation, modelInfo.c_str());
        return true;
    }
    // 在sdk层处理json字符串
    ai_runtime_core_text_embedding_service_complete_get_model_info(delegate, invocation, modelInfo.c_str());
    return true;
}

bool CoreTextEmbeddingService::parseTextEmbeddingInput(const std::string &jsonInput, std::string &text,
                                                       int &sessionId) const {
    Json::Value root;
    Json::Reader r;
    bool isReader = r.parse(jsonInput, root);
    if (!isReader) {
        LOGE("[parseTextEmbeddingInput] error: Failed to parse JSON input.");
        return false;
    }

    // 检查并提取 "text" 键
    if (root.isMember("text") && root["text"].isString()) {
        text = root["text"].asString();
    } else {
        LOGE("[parseTextEmbeddingInput] error: Key 'text' not found or not a string.");
        return false;
    }

    // 检查并提取 "sessionId" 键
    if (root.isMember("sessionId") && root["sessionId"].isInt()) {
        sessionId = root["sessionId"].asInt();
    } else {
        LOGE("[parseTextEmbeddingInput] error: Key 'sessionId' not found or not an integer.");
        return false;
    }

    return true;
}
