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

#include <iostream>

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

const char* imageEmbeddingProcessorObjectPath =
    "/com/kylin/AiRuntime/CoreImageEmbeddingService";
const std::string serverUnixPath = "unix:path=/tmp/.kylin-ai-runtime-unix/" +
                                   std::to_string(getuid()) +
                                   "/core-imageembedding.sock";
}  // namespace

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

int ImageEmbeddingProcessorProxy::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_image_embedding_service_proxy_new_sync(
        connection_, G_DBUS_PROXY_FLAGS_NONE, nullptr,
        imageEmbeddingProcessorObjectPath, nullptr, &error);
    if (delegate_ == nullptr && error != nullptr) {
        g_printerr("Error creating image embedding processor proxy %s: %s\n",
                   imageEmbeddingProcessorObjectPath, error->message);
        g_error_free(error);
        return COREAI_EMBEDDING_CONNECTION_ERROR;
    }
    return COREAI_EMBEDDING_SUCESS;
}

void ImageEmbeddingProcessorProxy::destroy() {
    if (connection_) {
        g_object_unref(connection_);
        connection_ = nullptr;
    }
    if (delegate_) {
        g_object_unref(delegate_);
        delegate_ = nullptr;
    }
}

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

    GError* error = nullptr;
    ai_runtime_core_image_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.\n", sessionId_);
    return COREAI_EMBEDDING_SUCESS;
}

EmbeddingResult* ImageEmbeddingProcessorProxy::embeddingText(
    const std::string& text) {
    GError* error = nullptr;
    nlohmann::json object;

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

    char* json_result = nullptr;

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

    if (error) {
        g_printerr("[Dbus-close:aisdk_image_embedding_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("Im-embedding text result error:json_result is nullptr!");
        EmbeddingResult* res = new EmbeddingResult(
            COREAI_EMBEDDING_RUNTIME_ERROR, "Embedding reulst is nullptr!");
        return res;
    }

    std::string string_json_result(json_result);
    g_free(json_result);

    auto result = parseResult(string_json_result);
    EmbeddingResult* res = new EmbeddingResult(result);
    return res;
}

EmbeddingResult* ImageEmbeddingProcessorProxy::embeddingImage(
    const std::string& file_path) {
    GError* error = nullptr;
    nlohmann::json object;

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

    char* json_result = nullptr;

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

    if (error) {
        g_printerr("[Dbus-close:aisdk_image_embedding_image]: %s\n",
                   error->message);
        // 由于runtime可能会崩溃，当前的session就不可用了，让用户重新创建sessino
        EmbeddingResult* res = new EmbeddingResult(
            COREAI_EMBEDDING_CONNECTION_ERROR, error->message);
        g_error_free(error);
        return res;
    }
    if (!json_result) {
        LOGE("Im-embedding image 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;
}

EmbeddingResult* ImageEmbeddingProcessorProxy::embeddingBase64Image(
    const std::string& file_data) {
    GError* error = nullptr;

    nlohmann::json object;

    object["base64ImageData"] = file_data;
    object["sessionId"] = sessionId_;

    char* json_result = nullptr;

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

    if (error) {
        g_printerr("[Dbus-close:aisdk_image_embedding_base64_image]: %s\n",
                   error->message);
        // 由于runtime可能会崩溃，当前的session就不可用了，让用户重新创建sessino
        EmbeddingResult* res = new EmbeddingResult(
            COREAI_EMBEDDING_CONNECTION_ERROR, error->message);
        g_error_free(error);
        return res;
    }
    if (!json_result) {
        LOGE("Im-embedding base64 image 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 ImageEmbeddingProcessorProxy::embeddingTextAsync(
    const std::string& text, ImageEmbeddingResultCallback 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_image_embedding_service_call_embedding_text(
        delegate_, object.dump().c_str(), nullptr,
        (GAsyncReadyCallback)embeddingTextAsyncCallback, context);
}

void ImageEmbeddingProcessorProxy::embeddingImageAsync(
    const std::string& file_path, ImageEmbeddingResultCallback callback,
    void* callback_user_data) {
    nlohmann::json object;

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

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

    ai_runtime_core_image_embedding_service_call_embedding_image_file(
        delegate_, object.dump().c_str(), nullptr,
        (GAsyncReadyCallback)embeddingImageAsyncCallback, context);
}

void ImageEmbeddingProcessorProxy::embeddingBase64ImageAsync(
    const std::string& file_data, ImageEmbeddingResultCallback callback,
    void* callback_user_data) {
    nlohmann::json object;

    object["base64ImageData"] = file_data;
    object["sessionId"] = sessionId_;

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

    ai_runtime_core_image_embedding_service_call_embedding_base64_image(
        delegate_, object.dump().c_str(), nullptr,
        (GAsyncReadyCallback)embeddingBase64ImageAsyncCallback, context);
}

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

    char* json_result = nullptr;

    ai_runtime_core_image_embedding_service_call_embedding_text_finish(
        (AiRuntimeCoreImageEmbeddingService*)object, &json_result, res, &error);

    if (error) {
        g_printerr("[Dbus-close:aisdk_image_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("Im-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;
}

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

    char* json_result = nullptr;

    ai_runtime_core_image_embedding_service_call_embedding_image_file_finish(
        (AiRuntimeCoreImageEmbeddingService*)object, &json_result, res, &error);

    if (error) {
        g_printerr("[Dbus-close:aisdk_image_embedding_image_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("Im-embedding image 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;
}

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

    char* json_result = nullptr;

    ai_runtime_core_image_embedding_service_call_embedding_base64_image_finish(
        (AiRuntimeCoreImageEmbeddingService*)object, &json_result, res, &error);

    if (error) {
        g_printerr(
            "[Dbus-close:aisdk_image_embedding_base64_image_async]: %s\n",
            error->message);
        // 由于runtime可能会崩溃，当前的session就不可用了，让用户重新创建sessino
        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(
            "Im-embedding base64 image 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;
}

EmbeddingResult ImageEmbeddingProcessorProxy::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;
}

std::string ImageEmbeddingProcessorProxy::getModelInfo() {
    GError* error = nullptr;
    char* model_info = nullptr;
    ai_runtime_core_image_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("image_model")) {
        LOGE(
            "Te-embedding parse modelInfo error! image_model not found in "
            "json.");
        g_free(model_info);
        return "";
    }

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

    return res;
}