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

#include "error.h"
#include "logger.h"
#include "nlohmann/json.hpp"

namespace {

const char* textEmbeddingProcessorObjectPath =
    "/com/kylin/AiRuntime/CoreTextEmbeddingService";

const std::string serverUnixPath = "unix:path=/tmp/.kylin-ai-runtime-unix/" +
                                   std::to_string(getuid()) +
                                   "/core-textembedding.sock";

}  // namespace

int TextEmbeddingProcessorProxy::initEngine() {
    nlohmann::json config;
    config["engineName"] = "Embedding";

    GError* error = nullptr;
    ai_runtime_core_text_embedding_service_call_init_sync(
        delegate_, config.dump().c_str(), &sessionId_, &initErrorCode_, nullptr,
        &error);

    if (error) {
        g_printerr("Init engine error: %s\n", error->message);
        g_error_free(error);
        return COREAI_EMBEDDING_CONNECTION_ERROR;
    }
    if (sessionId_ == -1) {
        LOGEF("Proxy init engine error sessionId:%i.", sessionId_);
        return COREAI_EMBEDDING_INIT_ERROR;
    }
    LOGIF("Proxy init engine sessionId:%i.", sessionId_);
    return COREAI_EMBEDDING_SUCESS;
}

int TextEmbeddingProcessorProxy::connect() {
    // 获得连接
    GError* error = nullptr;
    connection_ = g_dbus_connection_new_for_address_sync(
        serverUnixPath.c_str(), G_DBUS_CONNECTION_FLAGS_AUTHENTICATION_CLIENT,
        nullptr, /* GDBusAuthObserver */
        nullptr, /* GCancellable */
        &error);
    if (connection_ == nullptr && error != nullptr) {
        g_printerr("Error connecting to D-Bus address %s: %s\n",
                   serverUnixPath.c_str(), error->message);
        g_error_free(error);

        return COREAI_EMBEDDING_CONNECTION_ERROR;
    }

    // 获得代理
    delegate_ = ai_runtime_core_text_embedding_service_proxy_new_sync(
        connection_, G_DBUS_PROXY_FLAGS_NONE, nullptr,
        textEmbeddingProcessorObjectPath, nullptr, &error);
    if (delegate_ == nullptr && error != nullptr) {
        g_printerr("Error creating text embedding processor proxy %s: %s\n",
                   textEmbeddingProcessorObjectPath, error->message);
        g_error_free(error);
        return COREAI_EMBEDDING_CONNECTION_ERROR;
    }
    return COREAI_EMBEDDING_SUCESS;
}

TextEmbeddingProcessorProxy::~TextEmbeddingProcessorProxy() { destroy(); }

EmbeddingResult* TextEmbeddingProcessorProxy::embeddingText(
    const std::string& text) {
    GError* error = nullptr;

    nlohmann::json object;

    object["text"] = text;
    object["sessionId"] = sessionId_;

    char* json_result = nullptr;

    ai_runtime_core_text_embedding_service_call_embedding_text_sync(
        delegate_, object.dump().c_str(), &json_result, nullptr, &error);

    if (error) {
        g_printerr("[Dbus-close:aisdk_text]: %s\n", error->message);
        // 由于runtime可能会崩溃，当前的session就不可用了，让用户重新创建session

        EmbeddingResult* res = new EmbeddingResult(
            COREAI_EMBEDDING_CONNECTION_ERROR, error->message);
        g_error_free(error);
        return res;
    }
    if (!json_result) {
        LOGE("Te-embedding text result error:json_result is nullptr!");
        EmbeddingResult* res = new EmbeddingResult(
            COREAI_EMBEDDING_RUNTIME_ERROR, "Embedding reulst is nullptr!");
        return res;
    }
    auto result = parseResult(json_result);
    g_free(json_result);

    EmbeddingResult* res = new EmbeddingResult(result);
    return res;
}

void TextEmbeddingProcessorProxy::embeddingTextAsync(
    const std::string& text, TextEmbeddingResultCallback callback,
    void* callback_user_data) {
    nlohmann::json object;

    object["text"] = text;
    object["sessionId"] = sessionId_;

    EmbeddingRequestContext* context =
        new EmbeddingRequestContext{callback, callback_user_data};

    ai_runtime_core_text_embedding_service_call_embedding_text(
        delegate_, object.dump().c_str(), nullptr,
        (GAsyncReadyCallback)embeddingTextAsyncCallback, context);
}

void TextEmbeddingProcessorProxy::destroy() {
    // 如果用户想用新的连接重连，这里清理一下connection
    if (connection_) {
        g_object_unref(connection_);
        connection_ = nullptr;
    }

    if (delegate_) {
        g_object_unref(delegate_);
        delegate_ = nullptr;
    }
}

EmbeddingResult TextEmbeddingProcessorProxy::parseResult(
    const std::string& json_string_result) {
    EmbeddingResult result;

    nlohmann::json root =
        nlohmann::json::parse(json_string_result, nullptr, false);

    if (root.empty() || !root.contains("vector_result") ||
        !root["vector_result"].is_array() || !root.contains("errorMessage") ||
        !root.contains("errorCode")) {
        LOGE("Json format error.");
        return EmbeddingResult(COREAI_EMBEDDING_RUNTIME_ERROR,
                               "Json format error.");
    }

    auto vector_result = root["vector_result"].get<std::vector<float>>();
    for (const auto& value : vector_result) {
        result.vector_result.push_back(static_cast<float>(value));
    }

    result.error_message = root["errorMessage"].get<std::string>();
    result.error_code = root["errorCode"].get<int>();

    return result;
}

void TextEmbeddingProcessorProxy::embeddingTextAsyncCallback(GObject* object,
                                                             GAsyncResult* res,
                                                             gpointer data) {
    auto* context = static_cast<EmbeddingRequestContext*>(data);
    GError* error = NULL;

    char* json_result = nullptr;

    ai_runtime_core_text_embedding_service_call_embedding_text_finish(
        (AiRuntimeCoreTextEmbeddingService*)object, &json_result, res, &error);

    if (error) {
        g_printerr("[Dbus-close:aisdk_embedding_text_async]: %s\n",
                   error->message);

        // 由于runtime可能会崩溃，当前的session就不可用了，让用户重新创建session
        EmbeddingResult result(COREAI_EMBEDDING_CONNECTION_ERROR,
                               error->message);
        context->callback(&result, context->callback_user_data);
        g_error_free(error);
        delete context;
        context = nullptr;
        return;
    }
    if (!json_result) {
        LOGE(
            "Te-embedding text async result error:json_result is "
            "nullptr!");

        EmbeddingResult result(COREAI_EMBEDDING_RUNTIME_ERROR,
                               "Embedding reulst is nullptr!");
        context->callback(&result, context->callback_user_data);
        delete context;
        context = nullptr;
        return;
    }

    auto result = parseResult(json_result);
    g_free(json_result);
    context->callback(&result, context->callback_user_data);
    delete context;
    context = nullptr;
}

std::string TextEmbeddingProcessorProxy::getModelInfo() {
    GError* error = nullptr;
    char* model_info = nullptr;
    ai_runtime_core_text_embedding_service_call_get_model_info_sync(
        delegate_, sessionId_, &model_info, nullptr, &error);
    if (error) {
        g_printerr("[Dbus-close:aisdk_text]: %s\n", error->message);
        g_error_free(error);
        return "";
    }
    if (!model_info) {
        LOGE("Te-embedding modelInfo is empty!");
        return "";
    }
    // engine中返回的向量化模型信息包括文本和图像的，在sdk层拆开json字符串

    nlohmann::json root = nlohmann::json::parse(model_info, nullptr, false);

    if (root.empty() || !root.contains("models") ||
        !root["models"].contains("text_model")) {
        LOGE(
            "Te-embedding parse modelInfo error! text_model not found in "
            "json.");
        g_free(model_info);
        return "";
    }

    // 提取 text_model 的部分
    std::string res = root["models"]["text_model"].dump();
    g_free(model_info);

    return res;
}