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

#include <chrono>
#include <thread>

#include "_chatmodelconfig.h"
#include "_chatresult.h"
#include "logger.h"

namespace kyai::genai::nlp {

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

const char* nlpResultSignal = "ChatNlpResult";
const int chatResultCallbackTimeoutMs = 60 * 60 * 1000;  // 1 hour timeout

namespace {

ChatResult dissectChatResult(GVariant* parameters) {
    GVariantIter iter;
    g_variant_iter_init(&iter, parameters);

    GVariant* reason = g_variant_iter_next_value(&iter);
    GVariant* finishReason = g_variant_iter_next_value(&iter);
    GVariant* errorCode = g_variant_iter_next_value(&iter);
    GVariant* errorMessage = g_variant_iter_next_value(&iter);
    GVariant* isEnd = g_variant_iter_next_value(&iter);
    GVariant* reasoningContent = g_variant_iter_next_value(&iter);

    if (reason == nullptr or reasoningContent == nullptr or
        finishReason == nullptr or errorCode == nullptr or
        errorMessage == nullptr or isEnd == nullptr) {
        LOGE("chat result error: result is nullptr!");
        if (reason) g_variant_unref(reason);
        if (reasoningContent) g_variant_unref(reasoningContent);
        if (finishReason) g_variant_unref(finishReason);
        if (errorCode) g_variant_unref(errorCode);
        if (errorMessage) g_variant_unref(errorMessage);
        if (isEnd) g_variant_unref(isEnd);
        return {};
    }

    const char* resultReason =
        static_cast<const char*>(g_variant_get_data(reason));
    const char* resultReasoningContent =
        static_cast<const char*>(g_variant_get_data(reasoningContent));
    const char* resultFinishReason =
        static_cast<const char*>(g_variant_get_data(finishReason));
    const int* resultErrorCode =
        static_cast<const int*>(g_variant_get_data(errorCode));
    const char* resultErrorMessage =
        static_cast<const char*>(g_variant_get_data(errorMessage));
    const bool* resultIsEnd =
        static_cast<const bool*>(g_variant_get_data(isEnd));

    if (resultReason == nullptr or resultReasoningContent == nullptr or
        resultFinishReason == nullptr or resultErrorCode == nullptr or
        resultErrorMessage == nullptr or resultIsEnd == nullptr) {
        LOGE("chat result error: result is nullptr!");
        if (reason) g_variant_unref(reason);
        if (reasoningContent) g_variant_unref(reasoningContent);
        if (finishReason) g_variant_unref(finishReason);
        if (errorCode) g_variant_unref(errorCode);
        if (errorMessage) g_variant_unref(errorMessage);
        if (isEnd) g_variant_unref(isEnd);
        return {};
    }

    ChatResult result{resultReason,       resultReasoningContent,
                      resultFinishReason, *resultErrorCode,
                      resultErrorMessage, *resultIsEnd};

    g_variant_unref(reason);
    g_variant_unref(reasoningContent);
    g_variant_unref(finishReason);
    g_variant_unref(errorCode);
    g_variant_unref(errorMessage);
    g_variant_unref(isEnd);

    return result;
}

std::string jsonToString(const Json::Value& value) {
    Json::StreamWriterBuilder builder;
    builder["emitUTF8"] = true;
    return Json::writeString(builder, value);
}

Json::Value stringToJson(const std::string& string) {
    Json::Value value;
    if (Json::Reader reader; reader.parse(string, value)) {
        return value;
    } else {
        return Json::objectValue;
    }
}
}  // namespace

GenAiNlpServiceProxy::~GenAiNlpServiceProxy() {
    if (not GenAiNlpServer::getInstance().available()) {
        g_printerr(
            "Error creating core speech server proxy: Server proxy "
            "connection is unavailable.\n");
        return;
    }

    auto* connection = GenAiNlpServer::getInstance().getConnection();
    if (subscriptionId_ != 0) {
        // 取消订阅
        g_dbus_connection_signal_unsubscribe(connection, subscriptionId_);
    }

    removeCallbackTimerIfExists();
}

bool GenAiNlpServiceProxy::connectToServer() {
    GenAiNlpServer::getInstance().init();
    if (not GenAiNlpServer::getInstance().available()) {
        g_printerr(
            "Error creating core speech server proxy: Server proxy "
            "connection is unavailable.\n");
        return false;
    }

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

    GError* error{nullptr};
    auto* connection = GenAiNlpServer::getInstance().getConnection();
    proxy_ = ai_runtime_gen_ai_nlp_service_proxy_new_sync(
        connection, G_DBUS_PROXY_FLAGS_NONE, nullptr, objectPath, nullptr,
        &error);
    if (proxy_ == nullptr) {
        if (error) {
            g_printerr("Error creating text processor proxy %s: %s\n",
                       objectPath, error->message);
            g_error_free(error);
        }

        return false;
    }

    g_dbus_proxy_set_default_timeout(G_DBUS_PROXY(proxy_), 2 * 60 * 60 * 1000);

    return true;
}
int GenAiNlpServiceProxy::initEngine(const ChatModelConfig& chatModelConfig,
                                     int* sessionId) {
    Json::Value chatModelConfigJson = modelConfigToJson(chatModelConfig);
    std::string config = jsonToString(chatModelConfigJson);

    int errorCode = AISDK_NO_ERROR;
    GError* error = nullptr;
    bool success = ai_runtime_gen_ai_nlp_service_call_init_sync(
        proxy_, config.c_str(), sessionId, &errorCode, nullptr, &error);

    if (not success) {
        if (error) {
            LOGE("Error calling init: {}:{}", error->code, error->message);
            g_error_free(error);
            return AISDK_RUNTIME_ERROR;
        }
    }

    if (*sessionId != -1) {
        connectChatResultSignals(*sessionId);
        reconnectAttempts_ = 0;
    }
    return errorCode;
}

void GenAiNlpServiceProxy::destroyEngine(int sessionId) {
    ai_runtime_gen_ai_nlp_service_call_destroy_sync(proxy_, sessionId, nullptr,
                                                    nullptr);
    reconnectAttempts_ = 0;
}

void GenAiNlpServiceProxy::setChatResultCallback(
    const ChatResultCallback& callback, void* userData) {
    chatResultCallback_ = callback;
    userData_ = userData;
}

void GenAiNlpServiceProxy::setReconnectedCallback(
    const ReconnectedCallback& callback) {
    reconnectedCallback_ = std::move(callback);
}

void GenAiNlpServiceProxy::chat(int sessionId,
                                const ChatModelConfig& chatModelConfig,
                                const std::string& question) {
    if (question.empty()) {
        handleMessagesEmptyError();
        return;
    }
    reconnectAttempts_ = 0;
    currentSessionId_ = sessionId;

    // Start the timeout timer
    startCallbackTimer();
    Json::Value request = modelConfigToJson(chatModelConfig);
    request["question"] = question;
    std::string temp = jsonToString(request);
    ai_runtime_gen_ai_nlp_service_call_chat(
        proxy_, sessionId, temp.c_str(), nullptr, onChatFinishedCallback, this);
}

void GenAiNlpServiceProxy::handleErrorOccurred(int errorCode) {
    LOGE("Error occurred with code: {}", errorCode);

    removeCallbackTimerIfExists();

    switch (errorCode) {
        case G_IO_ERROR_CLOSED:
            handleServerClosed();
            break;
        case G_IO_ERROR_TIMED_OUT:
            handleChatTimeout();
            break;
        default:
            handleGeneralError();
            break;
    }
}

void GenAiNlpServiceProxy::handleServerClosed() {
    LOGE("Server closed unexpectedly.");
    if (chatResultCallback_) {
        ChatResult result{
            "",  "", "", AISDK_RUNTIME_ERROR, "Server closed unexpectedly.",
            true};
        chatResultCallback_(&result, userData_);
    }

    if (reconnectServer()) {
        if (reconnectedCallback_) {
            reconnectedCallback_();
        }
    }
}

void GenAiNlpServiceProxy::handleChatTimeout() {
    if (chatResultCallback_) {
        ChatResult result{
            "", "", "", AISDK_MODEL_RUN_TIME_OUT, "Run model time out.", true};
        chatResultCallback_(&result, userData_);
    }
}

void GenAiNlpServiceProxy::handleGeneralError() {
    if (chatResultCallback_) {
        ChatResult result{
            "", "", "", AISDK_RUNTIME_ERROR, "General error occurred.", true};
        chatResultCallback_(&result, userData_);
    }
}

void GenAiNlpServiceProxy::handleMessagesEmptyError() {
    if (chatResultCallback_) {
        ChatResult result{
            "", "", "", NLP_PARAM_ERROR, "The input message is empty.", true};
        chatResultCallback_(&result, userData_);
    };
}

bool GenAiNlpServiceProxy::reconnectServer() {
    if (reconnectAttempts_ >= maxReconnectAttempts) {
        LOGE("Max reconnect attempts reached");
        return false;
    }

    reconnectAttempts_++;
    LOGD("Reconnect attempt {}.", reconnectAttempts_);

    if (!connectToServer()) {
        LOGE("Failed to reconnect to server");
        std::this_thread::sleep_for(std::chrono::seconds(1));
        return reconnectServer();
    }

    LOGI("Successfully reconnected to server.");

    return true;
}

void GenAiNlpServiceProxy::stopChat(int sessionId) {
    ai_runtime_gen_ai_nlp_service_call_stop_chat_sync(proxy_, sessionId,
                                                      nullptr, nullptr);
}

void GenAiNlpServiceProxy::clearChatMessage(int sessionId) {
    ai_runtime_gen_ai_nlp_service_call_clear_chat_message_sync(
        proxy_, sessionId, nullptr, nullptr);
}

void GenAiNlpServiceProxy::setChatSystemPromptId(int sessionId,
                                                 PromptId promptId) {
    ai_runtime_gen_ai_nlp_service_call_set_chat_system_prompt_id_sync(
        proxy_, sessionId, promptId, nullptr, nullptr);
}

void GenAiNlpServiceProxy::addSystemMessage(int sessionId,
                                            const std::string& message) {
    ai_runtime_gen_ai_nlp_service_call_add_system_message_sync(
        proxy_, sessionId, message.c_str(), nullptr, nullptr);
}

void GenAiNlpServiceProxy::chatWithHistoryMessages(
    int sessionId, const ChatModelConfig& chatModelConfig,
    const std::string& messages) {
    if (messages.empty()) {
        handleMessagesEmptyError();
        return;
    }
    currentSessionId_ = sessionId;
    // Start the timeout timer
    startCallbackTimer();
    Json::Value request = stringToJson(messages);
    request["top_k"] = chatModelConfig.topK_;
    request["model_name"] = chatModelConfig.name_;
    request["deploy_type"] = getModelConfigType(chatModelConfig);
    std::string temp = jsonToString(request);
    LOGD(temp);

    ai_runtime_gen_ai_nlp_service_call_chat_with_history_messages(
        proxy_, sessionId, temp.c_str(), nullptr, onChatFinishedCallback, this);
}

void GenAiNlpServiceProxy::connectChatResultSignals(int sessionId) {
    auto* connection = GenAiNlpServer::getInstance().getConnection();

    std::string interfaceName = interface + std::to_string(sessionId);
    subscriptionId_ = g_dbus_connection_signal_subscribe(
        connection, nullptr, interfaceName.c_str(), nlpResultSignal, objectPath,
        nullptr, G_DBUS_SIGNAL_FLAGS_NONE, onChatResultCallback, this, nullptr);
}

void GenAiNlpServiceProxy::onChatResultCallback(
    GDBusConnection* connection, const gchar* senderName,
    const gchar* objectPath, const gchar* interfaceName,
    const gchar* signalName, GVariant* parameters, gpointer userData) {
    auto* proxy = static_cast<GenAiNlpServiceProxy*>(userData);

    // Reset the timeout timer since we received a callback
    proxy->resetCallbackTimer();

    if (proxy->chatResultCallback_ != nullptr) {
        auto result = dissectChatResult(parameters);
        proxy->chatResultCallback_(&result, proxy->userData_);
    } else {
        LOGE("chat result callback is not set!");
    }
}

void GenAiNlpServiceProxy::onChatFinishedCallback(GObject* sourceObject,
                                                  GAsyncResult* res,
                                                  gpointer userData) {
    GError* error{nullptr};
    auto* proxy = static_cast<GenAiNlpServiceProxy*>(userData);
    proxy->removeCallbackTimerIfExists();
    bool result = ai_runtime_gen_ai_nlp_service_call_chat_finish(
        (AiRuntimeGenAiNlpService*)sourceObject, res, &error);
    if (not result) {
        if (error == nullptr) {
            g_printerr("Error calling chat: Unknown error\n");
        } else {
            g_printerr("Error calling chat: %s\n", error->message);
            proxy->handleErrorOccurred(error->code);
            g_error_free(error);
        }
    }
}

void GenAiNlpServiceProxy::startCallbackTimer() {
    removeCallbackTimerIfExists();

    // Start a new timer
    timeoutId_ = g_timeout_add(
        chatResultCallbackTimeoutMs,
        [](gpointer data) -> gboolean {
            auto* proxy = static_cast<GenAiNlpServiceProxy*>(data);
            proxy->handleCallbackTimeout();
            proxy->timeoutId_ = 0;   // Reset the ID since the source is removed
            return G_SOURCE_REMOVE;  // Don't repeat the timeout
        },
        this);
}

void GenAiNlpServiceProxy::removeCallbackTimerIfExists() {
    if (timeoutId_ > 0) {
        g_source_remove(timeoutId_);
        timeoutId_ = 0;
    }
}

void GenAiNlpServiceProxy::resetCallbackTimer() {
    // Reset the timer by starting a new one
    startCallbackTimer();
}

void GenAiNlpServiceProxy::handleCallbackTimeout() {
    LOGE("NLP result callback timeout after 2 minutes");

    stopChat(currentSessionId_);

    if (chatResultCallback_) {
        ChatResult result{"",
                          "",
                          "",
                          AISDK_MODEL_RUN_TIME_OUT,
                          "NLP result callback timeout after 2 minutes",
                          true};
        chatResultCallback_(&result, userData_);
    }
}

}  // namespace kyai::genai::nlp
