#include "modelinference.h"

#include "assistant/handlers/requesthandler.h"
#include "assistant/message.h"
#include "modelconfig/modelstatustracker.h"
#include "threadpool/async.h"
#include "utils/logger.h"
#include "utils/util.h"

int ModelInference::getCurrentModelMaxTokens() {
    if (chatEngine_ == nullptr) {
        return 0;
    }
    std::string currentModelName = chatEngine_->currentModel();
    LOGI("Chat model name: {}", currentModelName);

    Json::Value modelInfoJson = util::stringToJson(chatEngine_->modelInfo());
    int maxContextTokens{};
    for (const auto& item : modelInfoJson["models"]) {
        if (item["model"].asString() == currentModelName) {
            maxContextTokens = item["maxTokens"].asInt();
            break;
        }
    }
    return maxContextTokens;
}

bool ModelInference::isCloudModel() {
    if (chatEngine_ != nullptr) {
        return chatEngine_->isCloud();
    }
    return false;
}

std::future<bool> ModelInference::callChatModelInference(const std::string& modelInput) {
    // TODO 函数过长，后续优化
    LOGD("callChatModelInference {}", modelInput);

    return cpr::async([this, modelInput] {
        if (chatEngine_ == nullptr) {
            LOGE("Base chat nlp engine is not set.");
            DialogueCallback::ChatData chatData;
            chatData.type_ = DialogueCallback::ChatErrorOccurred;
            chatData.text_ = "";
            chatData.errorCode_ = int(assistant::ErrorCode::NlpModelNotFound);
            chatData.errorMessage_ = "模型加载失败";
            chatData.modelName_ = "";
            callChatModelCallback(chatData);
            return false;
        }

        LOGI("Chat engine name: {}", chatEngine_->engineName());
        LOGI("Chat model name: {}", chatEngine_->currentModel());
        auto callback = [this](const ai_engine::gen_ai::nlp::TextGenerationResult& result) {
            DialogueCallback::ChatData chatData;
            if (result.success) {
                chatData.type_ = DialogueCallback::ChatNormal;
            } else {
                chatData.type_ = DialogueCallback::ChatErrorOccurred;
            }
            chatData.text_ = result.text;
            chatData.errorCode_ = result.error.getCode();
            chatData.errorMessage_ = result.error.getMessage();
            chatData.modelName_ = result.modelName;
            if (chatEngine_->isCloud()) {
                Json::Value modelInfoJson = util::stringToJson(chatEngine_->modelInfo());
                for (const auto& item : modelInfoJson["models"]) {
                    if (item["name"].asString() == chatData.modelName_) {
                        chatData.feedbackUrl_ = item["feedbackUrl"].asString();
                        break;
                    }
                }
            }
            callChatModelCallback(chatData);
        };

        if (!isInited_) {
            auto result = chatEngine_->init();
            if (!result.success) {
                int errorCode = int(assistant::Error::getErrorCodeFromNlpEngine(result.error.getCode()));
                if (errorCode == int(assistant::ErrorCode::NetError)) {
                    bool success = handleNetworkDisconnection(modelInput, callback);
                    if (success) {
                        return true;
                    }
                }
                DialogueCallback::ChatData chatData;
                chatData.type_ = DialogueCallback::ChatErrorOccurred;
                chatData.text_ = "";
                chatData.errorCode_ = errorCode;
                chatData.errorMessage_ = result.error.getMessage();
                chatData.modelName_ = "";
                callChatModelCallback(chatData);
                return false;
            }
            isInited_ = true;
            LOGI("Chat engine init success.");
        }

        chatEngine_->setTextGenerationResultCallback(callback);

        ModelStatusTracker statusTracker(chatEngine_->currentModel());
        auto result = chatEngine_->chat(modelInput);
        if (!result.success) {
            int errorCode = int(assistant::Error::getErrorCodeFromNlpEngine(result.error.getCode()));
            if (errorCode == int(assistant::ErrorCode::NetError)) {
                // 网络错误，在这里重选本地模型回答
                // 这里不报错，成功直接返回，不成功继续
                bool success = handleNetworkDisconnection(modelInput, callback);
                if (success) {
                    return true;
                }
            }
            DialogueCallback::ChatData chatData;
            chatData.type_ = DialogueCallback::ChatErrorOccurred;
            chatData.text_ = "";
            chatData.errorCode_ = errorCode;
            chatData.errorMessage_ = result.error.getMessage();
            chatData.modelName_ = "";

            callChatModelCallback(chatData);
            LOGD("nlp engine error {} {}", errorCode, result.error.getMessage());
            return false;
        }
        return true;
    });
}

DialogueData::CallSuccessIdentifier ModelInference::callIntentionModelInference(const std::string& modelInput) {
    LOGD("意图推理功能输入：{}", modelInput);

    DialogueData::CallSuccessIdentifier ret{};

    if (intentEngine_ == nullptr) {
        LOGE("意图识别引擎未设置");
        ret.successRes = false;
        ret.intentSupported = false;
        ret.errorCode = int(assistant::ErrorCode::IntentNlpModelNotFound);
        ret.errorMessage = "意图识别引擎未设置";
        return ret;
    }

    ret.intentSupported = true;
    intentEngine_->setTextGenerationResultCallback(
        [this](const ai_engine::gen_ai::nlp::TextGenerationResult& result) { handleChatResult(result); });

    Json::Value userInputJson = util::stringToJson(modelInput);
    userInputJson["task_id"] = 1;  // 调用端侧意图模型必加
    std::string userInput = util::jsonToString(userInputJson);

    LOGI("Chat engine name: {}", intentEngine_->engineName());

    ModelStatusTracker statusTracker(intentEngine_->currentModel());

    LOGD("userInput: {}", userInput);

    auto result = intentEngine_->chat(userInput);
    if (!result.success) {
        int errorCode = int(assistant::Error::getErrorCodeFromNlpEngine(result.error.getCode()));
        LOGD("nlp engine error {} {}", errorCode, result.error.getMessage());
        ret.errorCode = errorCode;
        ret.errorMessage = result.error.getMessage();
        return ret;
    }

    // 开始：处理前后 ```json```
    if (intentResult_.substr(0, 7) == "```json") {
        intentResult_ = intentResult_.substr(7);
    }
    if (intentResult_.substr(intentResult_.size() - 3, intentResult_.size()) == "```") {
        intentResult_ = intentResult_.substr(0, intentResult_.size() - 3);
    }
    // 结束：处理前后 ```json```

    ret.successRes = true;
    ret.responseContent = intentResult_;
    ret.modelName = intentModelName_;
    return ret;
}

std::string ModelInference::intentModelName() {
    if (intentEngine_ == nullptr) {
        return "";
    }
    return intentEngine_->currentModel();
}

void ModelInference::initChatEngine(const std::shared_ptr<assistant::TextGenerationEngine>& nlpEngine,
                                    EngineSelector& engineSelector) {
    chatEngine_ = nlpEngine;
    engineSelector_ = &engineSelector;  // TODO 暂时放这里赋值，后续需要优化
    isInited_ = false;
}

void ModelInference::initIntentEngine(const std::shared_ptr<assistant::TextGenerationEngine>& nlpEngine) {
    intentEngine_ = nlpEngine;
    if (intentEngine_ == nullptr) {
        // 意图引擎初始化时出错，不报错处理
        return;
    }
    auto result = intentEngine_->init();
    if (!result.success) {
        LOGI("Intent engine init failed, error info: {} {}.", result.error.getCode(), result.error.getMessage());
    }
}

void ModelInference::handleChatResult(const ai_engine::gen_ai::nlp::TextGenerationResult& result) {
    Json::Value resultText = util::stringToJson(result.text);
    // 处理reason
    if (resultText.isMember("result")) {
        intentReplyCache_ += resultText["result"].asString();
    }
    // 处理is_end
    if (resultText.isMember("is_end")) {
        if (resultText["is_end"].asBool()) {
            intentResult_ = intentReplyCache_;
            intentReplyCache_.clear();
            intentModelName_ = result.modelName;
        }
    }
}

bool ModelInference::handleNetworkDisconnection(
    const std::string& modelInput,
    const std::function<void(const ai_engine::gen_ai::nlp::TextGenerationResult& result)>& callback) {
    if (engineSelector_ == nullptr) {
        return false;
    }
    DeployType type{DeployType::Unknown};
    if (chatEngine_->isCustomModelEngine()) {
        type = DeployType::Custom;
    } else if (chatEngine_->isCloud()) {
        type = DeployType::PublicCloud;
    } else {
        type = DeployType::Local;
    }

    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;
    }
    engine->setTextGenerationResultCallback(callback);
    auto result = engine->chat(modelInput);
    return result.success;
}
