/*
 * 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 "chatbot.h"
//
#include <jsoncpp/json/reader.h>
#include <jsoncpp/json/value.h>
#include <modelconfig/modelconfigenums.h>

#include "config/engineconfiguration.h"
#include "error.h"
#include "modelconfig/modelstatustracker.h"
#include "taskmanager/taskrunner.h"
#include "utils/logger.h"
#include "utils/util.h"

Chatbot::Chatbot(EngineSelector& engineSelector, bool isDefaultConfig)
    : engineSelector_(engineSelector), isDefaultConfig_(isDefaultConfig) {}

Chatbot::~Chatbot() {
    destroy();
    LOGD("~Chatbot");
}

void Chatbot::chat(const std::string& request, ChatResultCallback chatResultCallback,
                   ChatFinishedCallback finishedCallback) {
    Json::Value requestBody = util::stringToJson(request);
    ChatRequest chatRequest{requestBody, false, chatResultCallback, finishedCallback};
    chatRequestQueue_.enqueue(chatRequest);
    chatAsync();
}

void Chatbot::chatWithHistoryMessage(const std::string& messages, ChatResultCallback chatResultCallback,
                                     ChatFinishedCallback finishedCallback) {
    Json::Value requestBody = util::stringToJson(messages);
    ChatRequest chatRequest{requestBody, true, chatResultCallback, finishedCallback};
    chatRequestQueue_.enqueue(chatRequest);
    chatAsync();
}

void Chatbot::stopChat() {
    clearChatContext();
    nlpEngine_->stopChat();
}
void Chatbot::clearChatContext() { chatMessages_.clear(); }

void Chatbot::addSystemMessage(const std::string& message) { chatMessages_.setSystemMessage(message, -1); }

void Chatbot::addSystemMessage(const std::string& message, int taskId) {
    chatMessages_.setSystemMessage(message, taskId);
}

void Chatbot::addSystemMessageByPromptId(int promptId, DeployType deployType) {
    using PromptId = config::PromptTemplate::PromptId;
    if (promptId == static_cast<int>(PromptId::DocTagExtraction) && deployType == DeployType::Local) {
        LOGI("Current promptId is {}, deployType is {}", promptId, int(deployType));
        LOGI("Current prompt content: {}", promptTemplate_->prompt(promptId).content);
        addSystemMessage(promptTemplate_->prompt(promptId).content, static_cast<int>(TaskId::TagExtraction));
        return;
    }
    addSystemMessage(promptTemplate_->prompt(promptId).content);
}

void Chatbot::addUserMessage(const std::string& message) { chatMessages_.addUserQuestion(message); }

void Chatbot::addAssistantMessage(const std::string& message) { chatMessages_.addAssistantReply(message); }

int Chatbot::init(const std::shared_ptr<INlpEngine>& nlpEngine) {
    try {
        if (nlpEngine == nullptr) {
            return AISDK_MODEL_NOT_FOUND;
        }
        nlpEngine_ = nlpEngine;
        promptTemplate_ = std::make_shared<config::PromptTemplate>(
            config::defaultPromptTemplateConfigFile(nlpEngine->engineName()),
            config::defaultPromptFilePath(nlpEngine->engineName(), nlpEngine_->currentModel()));

        auto result = nlpEngine_->init();
        if (!result.success) {
            LOGE("Failed to init nlp engine: {}", result.error.getMessage());
            return gen_ai::nlp::error::getNlpErrorCode(result.error.getCode());
            isInited_ = false;
        }
        isInited_ = true;
        std::shared_ptr sharedObject = getSharedObject();
        nlpEngine_->setTextGenerationResultCallback(
            [sharedObject](const ai_engine::gen_ai::nlp::TextGenerationResult& result) {
                sharedObject->handleChatResult(result);
            });

        return AISDK_NO_ERROR;
    } catch (std::exception& except) {
        LOGE("{}", except.what());
        return AISDK_RUNTIME_ERROR;
    }
}

void Chatbot::destroy() {
    try {
        stopChat();
        nlpEngine_->setTextGenerationResultCallback(nullptr);
    } catch (std::exception& except) {
        LOGE("{}", except.what());
    }
}

bool Chatbot::isDefaultConfig() { return isDefaultConfig_; }

void Chatbot::handleChatResult(const ai_engine::gen_ai::nlp::TextGenerationResult& result) {
    ChatResult chatResult;

    Json::Value resultText = util::stringToJson(result.text);
    // 处理reasoning_content
    if (resultText.isMember("reasoning_content")) {
        chatResult.reasoningContent_ = resultText["reasoning_content"].asString();
    }
    // 处理reason
    if (resultText.isMember("result")) {
        chatResult.reason_ = resultText["result"].asString();
        replyCache_ += chatResult.reason_;
    }
    // 处理is_end
    if (resultText.isMember("is_end")) {
        chatResult.isEnd_ = resultText["is_end"].asBool();
        if (chatResult.isEnd_ && !currentChatRequest_.withHistoryMessage) {
            chatMessages_.addAssistantReply(replyCache_);
            replyCache_.clear();
        }
    }
    // 处理finishReason
    if (resultText.isMember("finish_reason")) {
        chatResult.finishReason_ = resultText["finish_reason"].asString();
    }
    // 处理modelName
    chatResult.modelName_ = result.modelName;
    chatResult.error_ = result.error;
    callChatResultCallback(chatResult);
}

void Chatbot::callChatResultCallback(const ChatResult& result) const {
    if (currentChatRequest_.resultCallback != nullptr) {
        currentChatRequest_.resultCallback(result);
    }
}

void Chatbot::callChatFinishedCallback() const {
    if (currentChatRequest_.finishedCallback != nullptr) {
        currentChatRequest_.finishedCallback();
    }
}

void Chatbot::chatAsync() {
    if (taskId_ != task::INVALID_TASK_ID) {
        return;
    }
    taskId_ = task::generateTaskId();
    std::shared_ptr sharedObject = getSharedObject();
    task::TaskRunner::getInstance().addTask({taskId_, task::TaskType::TextChat, [sharedObject]() {
                                                 sharedObject->chatTask();
                                                 sharedObject->taskId_ = task::INVALID_TASK_ID;
                                             }});
}

void Chatbot::chatTask() {
    while (chatRequestQueue_.wait_dequeue_timed(currentChatRequest_, 10 * 1000 * 1000)) {
        std::string messages;
        if (currentChatRequest_.withHistoryMessage) {
            messages = util::jsonToString(currentChatRequest_.body);
        } else {
            std::string question = currentChatRequest_.body["question"].asString();
            chatMessages_.addUserQuestion(question);
            chatMessages_.setTopK(currentChatRequest_.body["top_k"].asInt());
            messages = chatMessages_.generateJsonString();
        }
        if (!isInited_) {
            promptTemplate_ = std::make_shared<config::PromptTemplate>(
                config::defaultPromptTemplateConfigFile(nlpEngine_->engineName()),
                config::defaultPromptFilePath(nlpEngine_->engineName(), nlpEngine_->currentModel()));
            auto result = nlpEngine_->init();
            if (!result.success) {
                LOGE("Failed to init nlp engine: {}", result.error.getMessage());
                int errorCode = gen_ai::nlp::error::getNlpErrorCode(result.error.getCode());
                if (errorCode == AISDK_NET_ERROR) {
                    bool success = handleNetworkDisconnection(messages);
                    if (success) {
                        callChatFinishedCallback();
                        return;
                    }
                }
                callChatResultCallback({"", "", "", "", result.error});
                return;
            }
            isInited_ = true;
            std::shared_ptr sharedObject = getSharedObject();
            nlpEngine_->setTextGenerationResultCallback(
                [sharedObject](const ai_engine::gen_ai::nlp::TextGenerationResult& result) {
                    sharedObject->handleChatResult(result);
                });
        }
        ModelStatusTracker statusTracker(nlpEngine_->currentModelName());
        auto result = nlpEngine_->chat(messages);
        if (!result.success) {
            int errorCode = gen_ai::nlp::error::getNlpErrorCode(result.error.getCode());
            if (errorCode == AISDK_NET_ERROR) {
                bool success = handleNetworkDisconnection(messages);
                if (success) {
                    callChatFinishedCallback();
                    return;
                }
            }
            callChatResultCallback({"", "", "", "", result.error});
        }
        callChatFinishedCallback();
    }
}

bool Chatbot::handleNetworkDisconnection(const std::string& message) {
    if (!isDefaultConfig_) {
        // 非系统默认配置条件，不触发断网使用本地逻辑
        return false;
    }
    DeployType type = currentEngineDeployType();

    EngineSelector::GenAiTextGenerationEnginePredicate predicate =
        [](const std::unique_ptr<ai_engine::gen_ai::nlp::AbstractTextGenerationEngine>& engine) {
            if (engine == nullptr) {
                return false;
            }
            auto result = engine->init();
            if (!result.success) {
                return false;
            }
            return true;
        };
    std::unique_ptr engine = engineSelector_.selectNextPriorityGenAiTextGenerationEngine(type, predicate);
    if (engine == nullptr) {
        return false;
    }
    std::shared_ptr sharedObject = getSharedObject();
    engine->setTextGenerationResultCallback([sharedObject](const ai_engine::gen_ai::nlp::TextGenerationResult& result) {
        sharedObject->handleChatResult(result);
    });
    auto result = engine->chat(message);
    return result.success;
}

DeployType Chatbot::currentEngineDeployType() const {
    DeployType type{DeployType::Unknown};
    if (nlpEngine_->isCustomModelEngine()) {
        type = DeployType::Custom;
    } else if (nlpEngine_->isCloud()) {
        type = DeployType::PublicCloud;
    } else {
        type = DeployType::Local;
    }
    return type;
}
