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

#include <kylin-ai/aiengine.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <sodium.h>

#include <cstdio>

#include "utils/logger.h"
#include "utils/util.h"

namespace {

constexpr const char* objectPath = "/com/kylin/AiRuntime/GenAiVision";
constexpr const char* interface = "com.kylin.AiRuntime.GenAiVision";
constexpr const char* visionResultSignal = "GenAiVisionResult";
constexpr const char* serviceName = "GenAiVisionService";

}  // namespace

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

GenAiVisionService::~GenAiVisionService() { unexportSkeleton(); }

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

void GenAiVisionService::exportSkeleton() {
    delegate_ = ai_runtime_gen_ai_vision_service_skeleton_new();
    connectSignal();

    GError* error = nullptr;

    isExported_ =
        g_dbus_interface_skeleton_export(G_DBUS_INTERFACE_SKELETON(delegate_), &connection_, objectPath, &error);
    if (!isExported_) {
        g_printerr("Error creating server at address %s: %s\n", objectPath, error->message);
        g_error_free(error);
    }
}

bool GenAiVisionService::onHandleInit(AiRuntimeGenAiVisionService* service, GDBusMethodInvocation* invocation,
                                      gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        LOGE("GenAiVisionService is null.");
        ai_runtime_gen_ai_vision_service_complete_init(service, invocation, -1, AISDK_RUNTIME_ERROR);
        return true;
    }

    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    auto visionEngine = sharedNlpService->engineSelector_.selectGenAiImageGenerationEngine();
    if (!visionEngine) {
        ai_runtime_gen_ai_vision_service_complete_init(service, invocation, -1, AISDK_MODEL_NOT_FOUND);
        LOGE("No gen ai vision engine found!");
        return true;
    }

    auto engineInitResult = visionEngine->init();
    if (!engineInitResult.success) {
        ai_runtime_gen_ai_vision_service_complete_init(service, invocation, -1,
                                                       genai::getVisionErrorCode(engineInitResult.error.getCode()));
        return true;
    }

    std::shared_ptr<ImageGenerator> imageGenerator = std::make_shared<ImageGenerator>(std::move(visionEngine));
    int sessionId = sharedNlpService->generateRandomSessionId();
    sharedNlpService->imageGenerator_[sessionId] = imageGenerator;

    std::string uuid = serviceName + std::to_string(sessionId);
    ModelConfigManager::getInstance().updateAllDeployTypeCallback(
        AiCapability::Vision, uuid,
        [sharedNlpService, sessionId] { sharedNlpService->modelChangedMap_[sessionId] = true; });

    ai_runtime_gen_ai_vision_service_complete_init(service, invocation, sessionId, AISDK_NO_ERROR);

    return true;
}

bool GenAiVisionService::onHandleGenerateImage(AiRuntimeGenAiVisionService* service, GDBusMethodInvocation* invocation,
                                               gint sessionId, const gchar* prompt, gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        g_dbus_method_invocation_return_error(invocation, G_IO_ERROR, G_IO_ERROR_FAILED,
                                              "Genai vision service not found.");
        LOGE("GenAiVisionService is null.");
        return false;
    }

    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    if (!sharedNlpService->isSessionIdValid(sessionId)) {
        g_dbus_method_invocation_return_error(invocation, G_IO_ERROR, G_IO_ERROR_FAILED, "No image generator found.");
        LOGE("No image generator found for session ID: {}", sessionId);
        return false;
    }

    sharedNlpService->tryToChangeModel(sessionId);

    const std::string resultSignalInterfaceName = interface + std::to_string((int)sessionId);
    auto sharedServiceObj = sharedNlpService->getSharedObject();
    sharedNlpService->imageGenerator_[sessionId]->generateImage(
        prompt,
        [resultSignalInterfaceName, sharedServiceObj](const ImageGenerationResult& result) {
            sharedServiceObj->emitImageResultSignal(resultSignalInterfaceName, result);
        },
        [service, invocation]() { ai_runtime_gen_ai_vision_service_complete_generate_image(service, invocation); });
    return true;
}

bool GenAiVisionService::onHandleStopImageGeneration(AiRuntimeGenAiVisionService* service,
                                                     GDBusMethodInvocation* invocation, gint sessionId,
                                                     gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        LOGE("GenAiVisionService is null.");
        return true;
    }
    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    if (!sharedNlpService->isSessionIdValid(sessionId)) {
        LOGE("No image generator found for session ID: {}", sessionId);
        return true;
    }

    sharedNlpService->imageGenerator_[sessionId]->stopImageGeneration();
    return true;
}

bool GenAiVisionService::onHandleGetSupportedImageNumber(AiRuntimeGenAiVisionService* service,
                                                         GDBusMethodInvocation* invocation, gint sessionId,
                                                         gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        LOGE("GenAiVisionService is null.");
        ai_runtime_gen_ai_vision_service_complete_get_supported_image_number(service, invocation, -1);
        return true;
    }
    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    if (!sharedNlpService->isSessionIdValid(sessionId)) {
        LOGE("No image generator found for session ID: {}", sessionId);
        ai_runtime_gen_ai_vision_service_complete_get_supported_image_number(service, invocation, -1);
        return true;
    }
    int supportedImageNumber = sharedNlpService->imageGenerator_.at(sessionId)->getSupportedImageNumber();
    ai_runtime_gen_ai_vision_service_complete_get_supported_image_number(service, invocation, supportedImageNumber);
    return true;
}

bool GenAiVisionService::onHandleGetSupportedImageSize(AiRuntimeGenAiVisionService* service,
                                                       GDBusMethodInvocation* invocation, gint sessionId,
                                                       gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        LOGE("GenAiVisionService is null.");
        ai_runtime_gen_ai_vision_service_complete_get_supported_image_size(service, invocation, "");
        return true;
    }
    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    if (!sharedNlpService->isSessionIdValid(sessionId)) {
        LOGE("No image generator found for session ID: {}", sessionId);
        ai_runtime_gen_ai_vision_service_complete_get_supported_image_size(service, invocation, "");
        return true;
    }

    std::string supportedImageSize = sharedNlpService->imageGenerator_.at(sessionId)->getSupportedImageSize();

    ai_runtime_gen_ai_vision_service_complete_get_supported_image_size(service, invocation, supportedImageSize.c_str());
    return true;
}

bool GenAiVisionService::onHandleGetSupportedImageStyle(AiRuntimeGenAiVisionService* service,
                                                        GDBusMethodInvocation* invocation, gint sessionId,
                                                        gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        LOGE("GenAiVisionService is null.");
        ai_runtime_gen_ai_vision_service_complete_get_supported_image_style(service, invocation, "");
        return true;
    }
    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    if (!sharedNlpService->isSessionIdValid(sessionId)) {
        LOGE("No image generator found for session ID: {}", sessionId);
        ai_runtime_gen_ai_vision_service_complete_get_supported_image_style(service, invocation, "");
        return true;
    }
    std::string supportedImageStyle = sharedNlpService->imageGenerator_.at(sessionId)->getSupportedImageStyle();

    ai_runtime_gen_ai_vision_service_complete_get_supported_image_style(service, invocation,
                                                                        supportedImageStyle.c_str());
    return true;
}

bool GenAiVisionService::onHandleSetImageConfig(AiRuntimeGenAiVisionService* service, GDBusMethodInvocation* invocation,
                                                gint sessionId, const gchar* imageConfig, gpointer userData) {
    auto* visionService = static_cast<GenAiVisionService*>(userData);
    if (!visionService) {
        LOGE("GenAiVisionService is null.");
        ai_runtime_gen_ai_vision_service_complete_set_image_config(service, invocation);
        return true;
    }
    std::shared_ptr sharedNlpService = visionService->getSharedObject();

    if (!sharedNlpService->isSessionIdValid(sessionId)) {
        LOGE("No image generator found for session ID: {}", sessionId);
        ai_runtime_gen_ai_vision_service_complete_set_image_config(service, invocation);
        return true;
    }

    genai::ImageConfig config{sharedNlpService->imageConfigFromJson(imageConfig)};
    sharedNlpService->imageGenerator_.at(sessionId)->setImageConfig(config);
    ai_runtime_gen_ai_vision_service_complete_set_image_config(service, invocation);
    return true;
}

void GenAiVisionService::connectSignal() {
    gulong handlerId{};
    handlerId = g_signal_connect(delegate_, "handle-init", G_CALLBACK(onHandleInit), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle-generate-image", G_CALLBACK(onHandleGenerateImage), this);
    handlerIds_.push_back(handlerId);

    handlerId =
        g_signal_connect(delegate_, "handle-stop-image-generation", G_CALLBACK(onHandleStopImageGeneration), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle_get_supported_image_number",
                                 G_CALLBACK(onHandleGetSupportedImageNumber), this);
    handlerIds_.push_back(handlerId);

    handlerId =
        g_signal_connect(delegate_, "handle_get_supported_image_size", G_CALLBACK(onHandleGetSupportedImageSize), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle_get_supported_image_style",
                                 G_CALLBACK(onHandleGetSupportedImageStyle), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle_set_image_config", G_CALLBACK(onHandleSetImageConfig), this);
    handlerIds_.push_back(handlerId);
}

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

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

int GenAiVisionService::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;
}

void GenAiVisionService::setImageGeneratorCallback(const std::shared_ptr<ImageGenerator>& imageGenerator,
                                                   int sessionId) {
    std::string resultSignalInterfaceName = interface + std::to_string((int)sessionId);
    std::shared_ptr sharedObject = getSharedObject();
    imageGenerator->setImageGeneratorCallback([sharedObject, resultSignalInterfaceName](ImageGenerationResult result) {
        sharedObject->emitImageResultSignal(resultSignalInterfaceName, result);
    });
}

bool GenAiVisionService::isSessionIdValid(int sessionId) {
    auto it = imageGenerator_.find(sessionId);
    if (it != imageGenerator_.end()) {
        return true;
    }
    return false;
}

GVariant* GenAiVisionService::buildGVariantFromImageResult(const ImageGenerationResult& result) {
    GVariantBuilder builder;
    g_variant_builder_init(&builder, G_VARIANT_TYPE_TUPLE);
    g_variant_builder_add(&builder, "s", result.imageBase64Data.c_str());
    g_variant_builder_add(&builder, "i", result.format);
    g_variant_builder_add(&builder, "i", result.height);
    g_variant_builder_add(&builder, "i", result.width);
    g_variant_builder_add(&builder, "i", result.total);
    g_variant_builder_add(&builder, "i", result.index);
    g_variant_builder_add(&builder, "i", result.errorCode);
    g_variant_builder_add(&builder, "s", result.errorMessage.c_str());

    GVariant* resultVariant = g_variant_builder_end(&builder);
    return resultVariant;
}

void GenAiVisionService::emitImageResultSignal(const std::string& interfaceName, const ImageGenerationResult& result) {
    GVariant* resultVariant = buildGVariantFromImageResult(result);
    g_dbus_connection_emit_signal(&connection_, interface, objectPath, interfaceName.c_str(), visionResultSignal,
                                  resultVariant, nullptr);
}

genai::ImageConfig GenAiVisionService::imageConfigFromJson(const std::string& jsonStr) {
    Json::Value root;
    Json::Reader reader;

    bool parsingSuccessful = reader.parse(jsonStr, root);
    if (!parsingSuccessful) {
        LOGE("Failed to parse JSON: {}", reader.getFormattedErrorMessages());
        genai::ImageConfig defaultConfig = {0, 0, 0, 0, ""};
        return defaultConfig;
    }

    genai::ImageConfig config;
    config.width = root["width"].asInt();
    config.height = root["height"].asInt();
    config.size = root["size"].asInt();
    config.number = root["number"].asInt();
    config.style = root["style"].asString();

    return config;
}

void GenAiVisionService::tryToChangeModel(int sessionId) {
    if (modelChangedMap_.find(sessionId) == modelChangedMap_.end()) {
        return;
    }
    if (!modelChangedMap_[sessionId]) {
        return;
    }

    if (imageGenerator_[sessionId]->engineInUse()) {
        return;
    }

    auto visionEngine = engineSelector_.selectGenAiImageGenerationEngine();
    if (!visionEngine) {
        LOGE("No gen ai vision engine found!");
        return;
    }
    imageGenerator_[sessionId]->setVisionEngine(std::move(visionEngine));
    modelChangedMap_[sessionId] = false;
}
