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

#include <gio/gio.h>
#include <jsoncpp/json/json.h>
#include <kylin-ai/ai-engine/imagegenerationengine.h>
#include <kylin-ai/ai-engine/textgenerationengine.h>

#include <cstdio>
#include <memory>
#include <utility>

#include "assistantserviceglue.h"
#include "config/prompttemplate.h"
#include "utils/logger.h"
#include "utils/util.h"

namespace {

constexpr const char *objectPath = "/com/kylin/AiRuntime/Assistant";
constexpr const char *serviceName = "AssistantService";

const std::set<int> aiassistantPromptIdSet = {
    int(config::PromptTemplate::PromptId::TextExpansion),                 // 1
    int(config::PromptTemplate::PromptId::TextPolishing),                 // 2
    int(config::PromptTemplate::PromptId::TranslateChineseToEnglish),     // 3
    int(config::PromptTemplate::PromptId::FileQuestionAnswering),         // 4
    int(config::PromptTemplate::PromptId::TodoListGeneration),            // 5
    int(config::PromptTemplate::PromptId::WeeklyReportGeneration),        // 6
    int(config::PromptTemplate::PromptId::MeetingInformationExtraction),  // 7
    int(config::PromptTemplate::PromptId::ErrorCorrection),               // 8
    int(config::PromptTemplate::PromptId::TextSummary),                   // 9
};

// 商汤-SenseChat5 ak sk
const char *engineConfig = R"(
    {
        "apiKey" : "C0D7765375D14DFBBF40363097B8D69D",
        "secretKey" : "B4044273AF8F4C69A08A75E47AEB005F"
    }
    )";

}  // namespace

bool AssistantService::onHandleInit(AiRuntimeAssistant *object, GDBusMethodInvocation *invocation, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (assistant == nullptr) {
        LOGE("Assistant::onHandleInit: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();

    LOGI("Assistant::onHandleInit...");

    int sessionId = sharedAssistantService->generateRandomSessionId();
    sharedAssistantService->dialogueManager_->addSession(sessionId);

    sharedAssistantService->initNlpEngines(sessionId);

    using PromptId = config::PromptTemplate::PromptId;
    sharedAssistantService->setNlpEnginePromptId(sessionId, int(PromptId::Default));

    assistant::Error visionError = sharedAssistantService->initVisionEngine(sessionId);
    if (visionError) {
        LOGE("Init vision engine error, error info: {} {}.", int(visionError.code), visionError.message);
    }

    visionError = sharedAssistantService->initImageProcessingEngine(sessionId, "");
    if (visionError) {
        LOGE("Init image processing engine error, error info: {} {}.", int(visionError.code), visionError.message);
    }

    sharedAssistantService->dialogueManager_->initIntentionRecognition(sessionId);

    ai_runtime_assistant_complete_init(object, invocation, sessionId, 0, "");
    LOGI("Assistant::onHandleInit: sessionId={}", sessionId);
    return true;
}

bool AssistantService::onHandleInitWithChatHistory(AiRuntimeAssistant *object, GDBusMethodInvocation *invocation,
                                                   const gchar *messages, gint promptId, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleInit: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();

    int sessionId = sharedAssistantService->generateRandomSessionId();
    sharedAssistantService->dialogueManager_->addSession(sessionId);
    sharedAssistantService->dialogueManager_->setChatContextMessages(sessionId, promptId, messages);

    auto nlpError = sharedAssistantService->initNlpEngines(sessionId);
    if (nlpError) {
        ai_runtime_assistant_complete_init_with_chat_history(object, invocation, sessionId, int(nlpError.code),
                                                             nlpError.message.c_str());
        return true;
    }

    if (promptId == -1) {
        promptId = int(config::PromptTemplate::PromptId::Default);
    }

    sharedAssistantService->setNlpEnginePromptId(sessionId, promptId);

    assistant::Error visionError = sharedAssistantService->initVisionEngine(sessionId);
    if (visionError) {
        LOGE("Init vision engine error, error info {} {}.", int(visionError.code), visionError.message);
    }

    visionError = sharedAssistantService->initImageProcessingEngine(sessionId, "");
    if (visionError) {
        LOGE("Init image processing engine error, error info {} {}.", int(visionError.code), visionError.message);
    }

    sharedAssistantService->dialogueManager_->initIntentionRecognition(sessionId);

    ai_runtime_assistant_complete_init_with_chat_history(object, invocation, sessionId, 0, "");
    LOGI("Assistant::onHandleInitWithChatHistory: sessionId={}", sessionId);
    return true;
}

bool AssistantService::onHandleSupportedFeatures(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                                 gint sessionId, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (assistant == nullptr) {
        LOGE("Assistant::onHandleSupportedFeatures: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();

    bool isRagSupported = RagServerProxy::instance().isValid();
    bool isImageProcessingSupported = sharedAssistantService->engineSelector_.isLocalCoreAiImageProcessEngineExist();
    int features =
        TextChat | TextToImage | (isImageProcessingSupported ? ImageProcessing : 0) | (isRagSupported ? DocumentQA : 0);
    LOGI("Supported features: {:b}.", features);
    ai_runtime_assistant_complete_supported_features(delegate, invocation, features);
    return true;
}

void AssistantService::tryToChangeNlpModel(int sessionId) {
    if (nlpModelChangedMap_.find(sessionId) == nlpModelChangedMap_.end()) {
        return;
    }
    LOGI("The base nlp model need changed flag: {}.", nlpModelChangedMap_[sessionId]);
    if (!nlpModelChangedMap_[sessionId]) {
        return;
    }
    std::unique_ptr baseNlpEngine = initBaseNlpEngine();
    if (baseNlpEngine == nullptr) {
        LOGI("Try to change base nlp model, the nlp engine is nullptr.");
    }
    // 即使是空的也设置，后续有保护
    dialogueManager_->setBaseNlpEngine(sessionId, std::move(baseNlpEngine), engineSelector_);
    nlpModelChangedMap_[sessionId] = false;
}

void AssistantService::tryToChangeIntentNlpModel(int sessionId) {
    if (intentNlpModelChangedMap_.find(sessionId) == intentNlpModelChangedMap_.end()) {
        return;
    }
    LOGI("The intent nlp model need changed flag: {}.", intentNlpModelChangedMap_[sessionId]);
    if (!intentNlpModelChangedMap_[sessionId]) {
        return;
    }
    std::unique_ptr intentNlpEngine = engineSelector_.selectGenAiTextGenerationEngine(DeployType::Local);
    if (intentNlpEngine == nullptr) {
        LOGI("Try to change intent nlp model, the nlp engine is nullptr.");
    }
    // 即使是空的也设置，后续有保护
    dialogueManager_->setIntentNlpEngine(sessionId, std::move(intentNlpEngine));
    intentNlpModelChangedMap_[sessionId] = false;
}

void AssistantService::tryToChangeVisionModel(int sessionId) {
    if (visionModelChangedMap_.find(sessionId) == visionModelChangedMap_.end()) {
        return;
    }
    LOGI("The text-to-image vision model need changed flag: {}.", visionModelChangedMap_[sessionId]);
    if (!visionModelChangedMap_[sessionId]) {
        return;
    }
    std::unique_ptr visionEngine = engineSelector_.selectGenAiImageGenerationEngine();
    if (visionEngine == nullptr) {
        LOGI("Try to change vision model, the vision engine is nullptr.");
    }
    // 即使是空的也设置，后续有保护
    dialogueManager_->setVisionEngine(sessionId, std::move(visionEngine));
    visionModelChangedMap_[sessionId] = false;
}

bool AssistantService::onHandleChat(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation, gchar *message,
                                    gint sessionId, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleInit: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();

    sharedAssistantService->tryToChangeNlpModel(sessionId);
    sharedAssistantService->tryToChangeIntentNlpModel(sessionId);
    sharedAssistantService->tryToChangeVisionModel(sessionId);

    LOGI("Assistant::onHandleChat: sessionId={}", sessionId);
    ai_runtime_assistant_complete_chat(delegate, invocation);
    sharedAssistantService->doHandleChat(sessionId, message);

    return true;
}

bool AssistantService::onHandleStopChat(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation, gint sessionId,
                                        gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleStopChat: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();

    sharedAssistantService->stopNlpEngine(sessionId);
    sharedAssistantService->stopVisionEngine(sessionId);

    return true;
}

bool AssistantService::onHandleSetPromptId(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                           gint sessionId, gint promptId, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleSetActorId: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();
    sharedAssistantService->setNlpEnginePromptId(sessionId, promptId);
    return true;
}

bool AssistantService::onHandleClearContext(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                            gint sessionId, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleClearContext: userData is null.");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();
    using PromptId = config::PromptTemplate::PromptId;
    sharedAssistantService->setNlpEnginePromptId(sessionId, int(PromptId::Default));
    sharedAssistantService->clearContext(sessionId);
    return true;
}

bool AssistantService::onHandlePrompts(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                       gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandlePrompts: userData is null.");
        ai_runtime_assistant_complete_prompts(delegate, invocation, "");
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();
    ai_runtime_assistant_complete_prompts(delegate, invocation, sharedAssistantService->getAllPrompts().c_str());
    return true;
}

bool AssistantService::onHandleCreatePrompt(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                            gchar *promptName, gchar *promptContent, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleCreatePrompt: userData is null.");
        ai_runtime_assistant_complete_create_prompt(delegate, invocation, -1);
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();
    int promptId{-1};
    if (sharedAssistantService->promptManager_ != nullptr) {
        promptId = sharedAssistantService->promptManager_->createPrompt(
            sharedAssistantService->aiEnginePluginManager_.allEngineName(), promptName, promptContent);
    }
    ai_runtime_assistant_complete_create_prompt(delegate, invocation, promptId);
    return true;
}

bool AssistantService::onHandleDeletePrompt(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                            gint promptId, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleDeletePrompt: userData is null.");
        ai_runtime_assistant_complete_delete_prompt(delegate, invocation, false);
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();
    bool success{false};
    if (sharedAssistantService->promptManager_ != nullptr) {
        success = sharedAssistantService->promptManager_->deletePrompt(
            sharedAssistantService->aiEnginePluginManager_.allEngineName(), promptId);
    }
    ai_runtime_assistant_complete_delete_prompt(delegate, invocation, success);
    return true;
}

bool AssistantService::onHandleUpdatePrompt(AiRuntimeAssistant *delegate, GDBusMethodInvocation *invocation,
                                            gint promptId, gchar *promptName, gchar *promptContent, gpointer userData) {
    auto *assistant = static_cast<AssistantService *>(userData);
    if (!assistant) {
        LOGE("Assistant::onHandleUpdatePrompt: userData is null.");
        ai_runtime_assistant_complete_update_prompt(delegate, invocation, false);
        return true;
    }
    std::shared_ptr sharedAssistantService = assistant->shared_from_this();
    bool success{false};
    if (sharedAssistantService->promptManager_ != nullptr) {
        success = sharedAssistantService->promptManager_->updatePrompt(promptId, promptName, promptContent);
    }
    ai_runtime_assistant_complete_update_prompt(delegate, invocation, success);
    return true;
}

assistant::ErrorCode AssistantService::doHandleChat(int sessionId, const std::string &message) {
    return dialogueManager_->handleRequest(sessionId, message);
}

void AssistantService::clearContext(int sessionId) { dialogueManager_->clearContext(sessionId); }

std::string AssistantService::getAllPrompts() {
    if (promptManager_ == nullptr) {
        return "";
    }
    auto prompts = promptManager_->prompts();
    bool ragIsValid = RagServerProxy::instance().isValid();
    if (!ragIsValid) {
        prompts.erase(std::remove_if(prompts.begin(), prompts.end(),
                                     [](const config::PromptTemplate::Prompt &prompt) {
                                         return prompt.id ==
                                                int(config::PromptTemplate::PromptId::FileQuestionAnswering);
                                     }),
                      prompts.end());
    }
    return aiassistantPrompts(prompts);
}

AssistantService::AssistantService(GDBusConnection &connection, ai_engine::AiEnginePluginManager &aiEnginePluginManager,
                                   EngineSelector &engineSelector)
    : connection_(connection),
      dialogueManager_(std::make_shared<DialogueManager>(connection)),
      aiEnginePluginManager_(aiEnginePluginManager),
      engineSelector_(engineSelector) {
    exportSkeleton();
    connectHandlers();
    initPromptManager();
}

AssistantService::~AssistantService() {
    for (const auto &sessionId : sessionIdSet_) {
        std::string uuid = serviceName + std::to_string(sessionId);
        engineSelector_.updateAllDeployTypeCallback(AiCapability::Nlp, uuid, nullptr);
        engineSelector_.updateAllDeployTypeCallback(AiCapability::Vision, uuid, nullptr);
    }
    LOGD("AssistantService::~AssistantService completed.");
}

void AssistantService::stop() {
    for (const auto &handlerId : handlerIds_) {
        g_signal_handler_disconnect(delegate_, handlerId);
    }
    handlerIds_.clear();
    for (const auto &sessionId : sessionIdSet_) {
        stopNlpEngine(sessionId);
        stopVisionEngine(sessionId);
    }
}

void AssistantService::exportSkeleton() {
    delegate_ = ai_runtime_assistant_skeleton_new();

    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);
    }
}

void AssistantService::connectHandlers() {
    gulong handlerId{};
    handlerId = g_signal_connect(delegate_, "handle-init", G_CALLBACK(onHandleInit), this);
    handlerIds_.push_back(handlerId);
    handlerId =
        g_signal_connect(delegate_, "handle_init_with_chat_history", G_CALLBACK(onHandleInitWithChatHistory), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-supported-features", G_CALLBACK(onHandleSupportedFeatures), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-chat", G_CALLBACK(onHandleChat), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-stop-chat", G_CALLBACK(onHandleStopChat), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-set-prompt-id", G_CALLBACK(onHandleSetPromptId), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-clear-context", G_CALLBACK(onHandleClearContext), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-prompts", G_CALLBACK(onHandlePrompts), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-create-prompt", G_CALLBACK(onHandleCreatePrompt), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-delete-prompt", G_CALLBACK(onHandleDeletePrompt), this);
    handlerIds_.push_back(handlerId);
    handlerId = g_signal_connect(delegate_, "handle-update-prompt", G_CALLBACK(onHandleUpdatePrompt), this);
    handlerIds_.push_back(handlerId);
}

void AssistantService::initPromptManager() {
    promptManager_ = std::make_unique<PromptManager>("baidu", "ernie-4.0-8k");
}

int AssistantService::generateRandomSessionId() {
    int sessionId = std::rand();
    while (sessionIdSet_.find(sessionId) != sessionIdSet_.end()) {
        sessionId = std::rand();
    }
    sessionIdSet_.insert(sessionId);
    return sessionId;
}

std::unique_ptr<AssistantService::INlpEngine> AssistantService::initBaseNlpEngine() {
    std::unique_ptr nlpEngine = engineSelector_.selectGenAiTextGenerationEngine();
    if (nlpEngine == nullptr) {
        return nullptr;
    }
    const std::string currentModel = nlpEngine->engineName() == "custom" ? "custom" : nlpEngine->currentModel();
    promptManager_ = std::make_unique<PromptManager>(nlpEngine->engineName(), currentModel);
    return nlpEngine;
}

assistant::Error AssistantService::initNlpEngines(int sessionId) {
    std::unique_ptr baseNlpEngine = initBaseNlpEngine();
    dialogueManager_->setBaseNlpEngine(sessionId, std::move(baseNlpEngine), engineSelector_);

    std::unique_ptr intentNlpEngine = engineSelector_.selectGenAiTextGenerationEngine(DeployType::Local);
    dialogueManager_->setIntentNlpEngine(sessionId, std::move(intentNlpEngine));

    // TODO 因为向量化引擎和意图NLP引擎绑定，暂时先放这里初始化
    //      后续如有必要，可以独立出来初始化
    // 配置文本向量化表示engine
    auto embeddingEngine = engineSelector_.selectCoreAiEmbeddingEngine("Embedding");
    if (embeddingEngine != nullptr) {
        dialogueManager_->setEmbeddingEngine(sessionId, std::move(embeddingEngine));
    } else {
        LOGE("Init embedding engine failed.");
    }

    // 初始化模型切换回调函数
    std::string uuid = serviceName + std::to_string(sessionId);
    engineSelector_.updateCallback(AiCapability::Nlp, DeployType::Custom, uuid,
                                   [this, sessionId] { nlpModelChangedMap_[sessionId] = true; });
    engineSelector_.updateCallback(AiCapability::Nlp, DeployType::PublicCloud, uuid,
                                   [this, sessionId] { nlpModelChangedMap_[sessionId] = true; });
    engineSelector_.updateCallback(AiCapability::Nlp, DeployType::Local, uuid, [this, sessionId] {
        nlpModelChangedMap_[sessionId] = true;
        intentNlpModelChangedMap_[sessionId] = true;
    });

    return {assistant::ErrorCode::Ok};
}

assistant::Error AssistantService::initVisionEngine(int sessionId) {
    std::unique_ptr visionEngine = engineSelector_.selectGenAiImageGenerationEngine();
    if (visionEngine == nullptr) {
        LOGI("Init vision engine, the vision engine is nullptr.");
    }
    // 即使是空的也设置，后续有保护
    dialogueManager_->setVisionEngine(sessionId, std::move(visionEngine));

    std::string uuid = serviceName + std::to_string(sessionId);
    engineSelector_.updateAllDeployTypeCallback(AiCapability::Vision, uuid,
                                                [this, sessionId] { visionModelChangedMap_[sessionId] = true; });

    return {assistant::ErrorCode::Ok};
}

assistant::Error AssistantService::initImageProcessingEngine(int sessionId, const std::string &config) {
    std::unique_ptr imageProcessingEngine = engineSelector_.selectCoreAiImageProcessEngine();
    if (imageProcessingEngine == nullptr) {
        LOGE("Create ImageProcessing engine failed.");
        return {assistant::ErrorCode::VisionModelNotFound, "Create ImageProcessing engine failed."};
    }
    dialogueManager_->setImageProcessingEngine(sessionId, std::move(imageProcessingEngine));
    return {assistant::ErrorCode::Ok};
}

void AssistantService::stopNlpEngine(int sessionId) {
    LOGI("stop chat");
    dialogueManager_->stopChat(sessionId);
}
void AssistantService::stopVisionEngine(int sessionId) { dialogueManager_->stopGenerateImage(sessionId); }

void AssistantService::setNlpEnginePromptId(int sessionId, int promptId) {
    if (promptManager_ == nullptr) {
        LOGE("Prompt manager is nullptr, set nlp engine prompt id failed.");
        return;
    }
    const config::PromptTemplate::Prompt &prompt = promptManager_->prompt(promptId);
    dialogueManager_->setChatSystemRole(sessionId, promptId, prompt.name, prompt.content);
    LOGD("Set nlp engine prompt id: promptId: {} prompt.name: {} prompt.content: {}.", promptId, prompt.name,
         prompt.content);
}

std::string AssistantService::aiassistantPrompts(const std::vector<config::PromptTemplate::Prompt> &prompts) {
    Json::Value jsonArray(Json::arrayValue);
    for (const auto &prompt : prompts) {
        if (aiassistantPromptIdSet.find(prompt.id) == aiassistantPromptIdSet.end() && prompt.id <= 100) {
            continue;
        }
        Json::Value json;
        json["id"] = prompt.id;
        json["name"] = prompt.name;

        if (prompt.type == config::PromptTemplate::PromptContentAccessType::NoAccess) {
            json["content"] = "";
        } else {
            json["content"] = prompt.content;
        }
        json["type"] = (int)prompt.type;
        jsonArray.append(json);

        if (prompt.id == (int)config::PromptTemplate::PromptId::TranslateChineseToEnglish) {
            if (engineSelector_.isLocalGenAiTextGenerationEngineExist()) {
                // 添加“系统控制”的提示词
                Json::Value json;
                json["id"] = int(config::PromptTemplate::PromptId::IntentionRecognition);
                json["name"] = "系统控制";
                json["content"] = "";
                json["type"] = config::PromptTemplate::PromptContentAccessType::NoAccess;
                jsonArray.append(json);
            }
        }
    }

    std::string ret = util::jsonToString(jsonArray);

    LOGI("Aiassistant prompts json: {}.", ret);

    return ret;
}
