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

#include <jsoncpp/json/json.h>
#include <kylin-ai/aiengine.h>
#include <kylin-ai/common/error.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <sodium.h>

#include <cstdio>

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

static const char *objectPath = "/com/kylin/AiRuntime/CoreVision";
static const char *interface = "com.kylin.AiRuntime.CoreVision";

static const char *textRecognitionResultSignal = "TextRecognitionVisionResult";

namespace {
GVariant *buildGVariantFromTextRecognitionResult(const TextRecognitionResult &result) {
    GVariantBuilder resultBuilder;
    g_variant_builder_init(&resultBuilder, G_VARIANT_TYPE("(sa(sa(ii)d)is)"));

    // alltext
    g_variant_builder_add(&resultBuilder, "s", result.allText.c_str());

    // lines
    GVariantBuilder lines_builder;
    g_variant_builder_init(&lines_builder, G_VARIANT_TYPE("a(sa(ii)d)"));

    for (const auto &[text, points, conf] : result.lines) {
        // lines.text
        GVariantBuilder line_builder;
        g_variant_builder_init(&line_builder, G_VARIANT_TYPE("(sa(ii)d)"));
        g_variant_builder_add(&line_builder, "s", text.c_str());

        // lines.points
        GVariantBuilder points_builder;
        g_variant_builder_init(&points_builder, G_VARIANT_TYPE("a(ii)"));
        for (const PixelPoint &point : points) {
            g_variant_builder_add(&points_builder, "(ii)", point.x, point.y);
        }
        g_variant_builder_add_value(&line_builder, g_variant_builder_end(&points_builder));

        // lines.conf
        g_variant_builder_add(&line_builder, "d", static_cast<double>(conf));
        g_variant_builder_add_value(&lines_builder, g_variant_builder_end(&line_builder));
    }
    g_variant_builder_add_value(&resultBuilder, g_variant_builder_end(&lines_builder));

    g_variant_builder_add(&resultBuilder, "i", result.errorCode);
    g_variant_builder_add(&resultBuilder, "s", result.errorMessage.c_str());

    GVariant *resultVariant = g_variant_builder_end(&resultBuilder);

    return resultVariant;
}
}  // namespace

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

CoreVisionService::~CoreVisionService() {
    for (const auto &[sessionId, textRecognition] : textRecognitions_) {
        if (textRecognition != nullptr) {
            textRecognition->destroy();
        }
    }

    unexportSkeleton();
    LOGD("~CoreVisionService");
}

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

void CoreVisionService::exportSkeleton() {
    delegate_ = ai_runtime_core_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 CoreVisionService::onHandleInit(AiRuntimeCoreVisionService *service, GDBusMethodInvocation *invocation,
                                     const gchar *modelconfig, gpointer userData) {
    auto *visionService = static_cast<CoreVisionService *>(userData);
    if (!visionService) {
        LOGE("CoreVisionService is null.");
        ai_runtime_core_vision_service_complete_init(service, invocation, -1, AISDK_RUNTIME_ERROR);
        return true;
    }
    std::shared_ptr sharedVisionService = visionService->getSharedObject();

    auto visionEngine = sharedVisionService->selectTextRecognitionEngine();

    if (!visionEngine) {
        LOGE("No vision engine found!");
        ai_runtime_core_vision_service_complete_init(service, invocation, -1, AISDK_MODEL_NOT_FOUND);
        return true;
    }

    int sessionId = sharedVisionService->generateRandomSessionId();

    LOGI("init engine {}", visionEngine->engineName(), sessionId);
    auto textRecognition = std::make_shared<TextRecognition>(std::move(visionEngine));

    if (not textRecognition->init()) {
        sharedVisionService->releaseRandomSessionId(sessionId);
        ai_runtime_core_vision_service_complete_init(service, invocation, -1, AISDK_MODEL_RUN_FAILED);
        return true;
    }
    sharedVisionService->textRecognitions_[sessionId] = textRecognition;

    ai_runtime_core_vision_service_complete_init(service, invocation, sessionId, AISDK_NO_ERROR);

    return true;
}

bool CoreVisionService::onHandleDestroyTextRecognition(AiRuntimeCoreVisionService *service,
                                                       GDBusMethodInvocation *invocation, gint sessionId,
                                                       gpointer userData) {
    LOGD("onHandleDestroyTextRecognition");
    auto *visionService = static_cast<CoreVisionService *>(userData);
    if (visionService == nullptr) {
        LOGE("CoreVisionService is null.");
        return true;
    }
    std::shared_ptr sharedVisionService = visionService->getSharedObject();

    if (not sharedVisionService->isTextRecognitionSessionIdValid(sessionId)) {
        LOGE("No text recognition found for session ID: {}", sessionId);
        return true;
    }

    if (sharedVisionService->textRecognitions_.find(sessionId) != sharedVisionService->textRecognitions_.end()) {
        sharedVisionService->textRecognitions_[sessionId]->destroy();
    }
    sharedVisionService->textRecognitions_.erase(sessionId);

    ai_runtime_core_vision_service_complete_destroy(service, invocation);
    LOGD("onHandleDestroyTextRecognition complete");
    return true;
}

bool CoreVisionService::onHandleTextRecognitionFromImageData(AiRuntimeCoreVisionService *service,
                                                             GDBusMethodInvocation *invocation, int sessionId,
                                                             const gchar *imageData, unsigned int imageDataLength,
                                                             gpointer userData) {
    // TODO 因为传递的是base64数据，所以不需要imageDataLength
    auto *visionService = static_cast<CoreVisionService *>(userData);
    if (!visionService) {
        TextRecognitionResult result;
        result.errorCode = AISDK_RUNTIME_ERROR;
        result.errorMessage = "No text recognition service found.";
        completeRecognizeTextFromImageData(result, service, invocation);
        LOGE("CoreVisionService is null.");
        return true;
    }
    std::shared_ptr sharedVisionService = visionService->getSharedObject();

    if (!sharedVisionService->isTextRecognitionSessionIdValid(sessionId)) {
        TextRecognitionResult result;
        result.errorCode = AISDK_RUNTIME_ERROR;
        result.errorMessage = "No text recognition found.";
        completeRecognizeTextFromImageData(result, service, invocation);
        LOGE("No text recognition found for session ID: {}", sessionId);
        return true;
    }

    sharedVisionService->textRecognitions_[sessionId]->recognizeTextFromImageData(
        imageData, [service, invocation](const TextRecognitionResult &result) {
            completeRecognizeTextFromImageData(result, service, invocation);
        });

    LOGI("start text recognition.");
    return true;
}

bool CoreVisionService::onHandleTextRecognitionFromImageFile(AiRuntimeCoreVisionService *service,
                                                             GDBusMethodInvocation *invocation, int sessionId,
                                                             const gchar *imageFilePath, gpointer userData) {
    auto *visionService = static_cast<CoreVisionService *>(userData);
    if (!visionService) {
        TextRecognitionResult result;
        result.errorCode = AISDK_RUNTIME_ERROR;
        result.errorMessage = "No text recognition service found.";
        completeRecognizeTextFromImageFile(result, service, invocation);
        LOGE("CoreVisionService is null.");
        return true;
    }
    std::shared_ptr sharedVisionService = visionService->getSharedObject();

    if (!sharedVisionService->isTextRecognitionSessionIdValid(sessionId)) {
        TextRecognitionResult result;
        result.errorCode = AISDK_RUNTIME_ERROR;
        result.errorMessage = "No text recognition found.";
        completeRecognizeTextFromImageFile(result, service, invocation);
        LOGE("No text recognition found for session ID: {}", sessionId);
        return true;
    }

    sharedVisionService->textRecognitions_[sessionId]->recognizeTextFromImageFile(
        imageFilePath, [service, invocation](const TextRecognitionResult &result) {
            completeRecognizeTextFromImageFile(result, service, invocation);
        });

    return true;
}

void CoreVisionService::completeRecognizeTextFromImageFile(const TextRecognitionResult &result,
                                                           AiRuntimeCoreVisionService *service,
                                                           GDBusMethodInvocation *invocation) {
    GVariant *resultVariant = buildGVariantFromTextRecognitionResult(result);
    ai_runtime_core_vision_service_complete_recognize_text_from_image_file(service, invocation, resultVariant);
}

void CoreVisionService::completeRecognizeTextFromImageData(const TextRecognitionResult &result,
                                                           AiRuntimeCoreVisionService *service,
                                                           GDBusMethodInvocation *invocation) {
    GVariant *resultVariant = buildGVariantFromTextRecognitionResult(result);
    ai_runtime_core_vision_service_complete_recognize_text_from_image_data(service, invocation, resultVariant);
}

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

    handlerId = g_signal_connect(delegate_, "handle-destroy", G_CALLBACK(onHandleDestroyTextRecognition), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle-recognize-text-from-image-file",
                                 G_CALLBACK(onHandleTextRecognitionFromImageFile), this);
    handlerIds_.push_back(handlerId);

    handlerId = g_signal_connect(delegate_, "handle-recognize-text-from-image-data",
                                 G_CALLBACK(onHandleTextRecognitionFromImageData), this);
    handlerIds_.push_back(handlerId);
}

void CoreVisionService::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 CoreVisionService::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 CoreVisionService::releaseRandomSessionId(int sessionId) { sessionIdSet_.erase(sessionId); }

bool CoreVisionService::isTextRecognitionSessionIdValid(int sessionId) {
    return textRecognitions_.find(sessionId) != textRecognitions_.end();
}

DeployType CoreVisionService::parseDeployType(const std::string &typeStr) {
    if (typeStr == "Local") {
        return DeployType::Local;
    }
    if (typeStr == "PrivateCloud") {
        return DeployType::Custom;
    }
    if (typeStr == "PublicCloud") {
        return DeployType::PublicCloud;
    }
    return DeployType::Default;  // default
}

const std::unique_ptr<ai_engine::core_ai::vision::AbstractTextRecognitionEngine>
CoreVisionService::selectTextRecognitionEngine() {
    return engineSelector_.selectTextRecognitionEngine();
}