/*
 * 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 "coreimageembeddingservice.h"

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

#include "utils/logger.h"

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

std::set<int> CoreImageEmbeddingService::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

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

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

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

void CoreImageEmbeddingService::exportSkeleton() {
    delegate_ = ai_runtime_core_image_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-embedding-image-file", G_CALLBACK(onHandleEmbeddingImageFile), this);
    handlerIds_.push_back(handlerId);

    handlerId =
        g_signal_connect(delegate_, "handle-embedding-base64-image", G_CALLBACK(onHandleEmbeddingBase64Image), 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 CoreImageEmbeddingService::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 CoreImageEmbeddingService::onHandleInitEngine(AiRuntimeCoreImageEmbeddingService *delegate,
                                                   GDBusMethodInvocation *invocation, const gchar *inputJson,
                                                   gpointer userData) {
    std::shared_ptr<CoreImageEmbeddingService> imageEmbeddingService =
        static_cast<CoreImageEmbeddingService *>(userData)->getSharedObject();
    if (!imageEmbeddingService) {
        LOGE("[onHandleInitEngine] error: imageEmbeddingService is nullptr!");
        ai_runtime_core_image_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);
    const std::string key = "engineName";

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

    int sessionId = imageEmbeddingService->generateRandomSessionId();

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

    if (!initResult) {
        ai_runtime_core_image_embedding_service_complete_init(delegate, invocation, -1, -1);
        return true;
    }
    LOGI("[onHandleInitEngine] service: The sessionId of the engine is {}", sessionId);
    LOGI("[onHandleInitEngine] start load image embedding model.");
    ai_engine::result::AiEngineBaseResult res =
        imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->initImageEmbeddingModule();
    if (!res.success) {
        LOGE("[onHandleInitEngine] load image embedding model error: {}", res.error.getMessage());
        ai_runtime_core_image_embedding_service_complete_init(delegate, invocation, sessionId, -1);
        return true;
    }
    LOGI("[onHandleInitEngine] load image embedding model success.");

    int initModuleErrorCode = 0;
    ai_runtime_core_image_embedding_service_complete_init(delegate, invocation, sessionId, initModuleErrorCode);
    return true;
}

int CoreImageEmbeddingService::generateRandomSessionId() const {
    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 CoreImageEmbeddingService::initEngine(const std::string &engineName, gint sessionId) {
    auto imageEmbeddingEngine = engineSelector_.selectCoreAiEmbeddingEngine(engineName);
    if (!imageEmbeddingEngine) {
        LOGE("[initEngine] error: imageEmbeddingEngine is nullptr!");
        return false;
    }
    imageEmbeddingEngineMap_[sessionId] = std::move(imageEmbeddingEngine);

    return true;
}

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

        std::string response =
            createResponse({}, -1, "[onHandleEmbeddingText] error: imageEmbeddingService is nullptr!");
        ai_runtime_core_image_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 (!imageEmbeddingService->parseTextEmbeddingInput(input, text, sessionId)) {
        LOGE("[onHandleEmbeddingText] error: parse input json failed!");

        std::string response = createResponse({}, -1, "[onHandleEmbeddingText] error: parse input json failed!");
        ai_runtime_core_image_embedding_service_complete_embedding_text(delegate, invocation, response.c_str());
        return true;
    }

    if (imageEmbeddingService->imageEmbeddingEngineMap_.find(sessionId) ==
        imageEmbeddingService->imageEmbeddingEngineMap_.end()) {
        LOGE("[onHandleEmbeddingText] error: ImageEmbeddingEngine is nullptr!");

        std::string response =
            createResponse({}, -1, "[onHandleEmbeddingText] error: ImageEmbeddingEngine is nullptr!");

        ai_runtime_core_image_embedding_service_complete_embedding_text(delegate, invocation, response.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 response = createResponse({}, -1,
                                                  "[onHandleEmbeddingText] embedding engine "
                                                  "error: errorMessage: " +
                                                      errorMessage);
            ai_runtime_core_image_embedding_service_complete_embedding_text(delegate, invocation, response.c_str());
            return;
        }
        std::string response = createResponse(result.embeddings, 0, "Success");
        ai_runtime_core_image_embedding_service_complete_embedding_text(delegate, invocation, response.c_str());
    };

    // 目前is_question模型暂未支持
    imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->setEmbeddingResultCallback(callback);
    imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->embeddingTextWithImageModel(text);

    return true;
}
bool CoreImageEmbeddingService::onHandleEmbeddingImageFile(AiRuntimeCoreImageEmbeddingService *delegate,
                                                           GDBusMethodInvocation *invocation, const gchar *inputJson,
                                                           gpointer userData) {
    std::shared_ptr<CoreImageEmbeddingService> imageEmbeddingService =
        static_cast<CoreImageEmbeddingService *>(userData)->getSharedObject();
    if (!imageEmbeddingService) {
        LOGE("[onHandleEmbeddingImageFile] error: imageEmbeddingService is nullptr!");
        std::string response = createResponse({}, -1,
                                              "[onHandleEmbeddingImageFile] error: "
                                              "imageEmbeddingService is nullptr!");
        ai_runtime_core_image_embedding_service_complete_embedding_image_file(delegate, invocation, response.c_str());
        return true;
    }
    // parse input json
    std::string input(inputJson);
    std::string filePath;
    int sessionId;
    if (!imageEmbeddingService->parseImageEmbeddingInput(input, filePath, sessionId)) {
        LOGE("[onHandleEmbeddingImageFile] error: parse input json failed!");
        std::string response = createResponse({}, -1, "[onHandleEmbeddingImageFile] error: parse input json failed!");
        ai_runtime_core_image_embedding_service_complete_embedding_image_file(delegate, invocation, response.c_str());
        return true;
    }

    if (imageEmbeddingService->imageEmbeddingEngineMap_.find(sessionId) ==
        imageEmbeddingService->imageEmbeddingEngineMap_.end()) {
        LOGE("[onHandleEmbeddingImageFile] error: ImageEmbeddingEngine is nullptr!");

        std::string response = createResponse({}, -1,
                                              "[onHandleEmbeddingImageFile] error: "
                                              "ImageEmbeddingEngine is nullptr!");
        ai_runtime_core_image_embedding_service_complete_embedding_image_file(delegate, invocation, response.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("[onHandleEmbeddingImageFile] embedding engine error: errorMessage : {}", errorMessage);
            std::string response = createResponse({}, -1,
                                                  "[onHandleEmbeddingImageFile] embedding engine "
                                                  "error: errorMessage: " +
                                                      errorMessage);
            ai_runtime_core_image_embedding_service_complete_embedding_image_file(delegate, invocation,
                                                                                  response.c_str());
            return;
        }
        std::string response = createResponse(result.embeddings, 0, "Success");

        ai_runtime_core_image_embedding_service_complete_embedding_image_file(delegate, invocation, response.c_str());
    };
    imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->setEmbeddingResultCallback(callback);
    imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->embeddingImageFromImageFile(filePath);

    return true;
}

bool CoreImageEmbeddingService::onHandleEmbeddingBase64Image(AiRuntimeCoreImageEmbeddingService *delegate,
                                                             GDBusMethodInvocation *invocation, const gchar *inputJson,
                                                             gpointer userData) {
    std::shared_ptr<CoreImageEmbeddingService> imageEmbeddingService =
        static_cast<CoreImageEmbeddingService *>(userData)->getSharedObject();
    if (!imageEmbeddingService) {
        LOGE("[onHandleEmbeddingBase64Image] error: imageEmbeddingService is nullptr!");
        std::string response = createResponse({}, -1,
                                              "[onHandleEmbeddingBase64Image] error: "
                                              "imageEmbeddingService is nullptr!");
        ai_runtime_core_image_embedding_service_complete_embedding_base64_image(delegate, invocation, response.c_str());
        return true;
    }
    // parse input json
    std::string input(inputJson);
    std::string base64ImageData;
    int sessionId;
    if (!imageEmbeddingService->parseBase64ImageEmbeddingInput(input, base64ImageData, sessionId)) {
        LOGE("[onHandleEmbeddingBase64Image] error: parse input json failed!");

        std::string resultJson =
            createResponse({}, -1, "[onHandleEmbeddingBase64Image]error: parse input json failed!");

        ai_runtime_core_image_embedding_service_complete_embedding_base64_image(delegate, invocation,
                                                                                resultJson.c_str());
        return true;
    }

    if (imageEmbeddingService->imageEmbeddingEngineMap_.find(sessionId) ==
        imageEmbeddingService->imageEmbeddingEngineMap_.end()) {
        LOGE("[onHandleEmbeddingBase64Image] error: ImageEmbeddingEngine is nullptr!");

        std::string resultJson = createResponse({}, -1,
                                                "[onHandleEmbeddingBase64Image] error: "
                                                "ImageEmbeddingEngine is nullptr!");
        ai_runtime_core_image_embedding_service_complete_embedding_base64_image(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("[onHandleEmbeddingBase64Image] embedding engine error: errorMessage : {}", errorMessage);

            std::string resultJson = createResponse({}, -1,
                                                    "[onHandleEmbeddingBase64Image] embedding "
                                                    "engine error: errorMessage: " +
                                                        errorMessage);
            ai_runtime_core_image_embedding_service_complete_embedding_base64_image(delegate, invocation,
                                                                                    resultJson.c_str());
            return;
        }
        std::string resultJson = createResponse(result.embeddings, 0, "Success");

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

    imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->setEmbeddingResultCallback(callback);
    imageEmbeddingService->imageEmbeddingEngineMap_[sessionId]->embeddingImageFromBase64ImageData(base64ImageData);
    return true;
}

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

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

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

bool CoreImageEmbeddingService::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;
}
bool CoreImageEmbeddingService::parseImageEmbeddingInput(const std::string &jsonInput, std::string &filePath,
                                                         int &sessionId) const {
    Json::Value root;
    Json::Reader r;
    bool isReader = r.parse(jsonInput, root);
    if (!isReader) {
        LOGE("[parseImageEmbeddingInput] error: Failed to parse JSON input.");
        return false;
    }

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

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

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

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

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

    return true;
}
