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

#include <jsoncpp/json/json.h>

#include "utils/logger.h"
namespace {
Json::Value authensToEngineConfigJson(const std::vector<MultiAuthentication> &multiAuthens) {
    Json::Value root;
    for (const auto &multiAuth : multiAuthens) {
        for (const auto &auth : multiAuth.authentications) {
            root[auth.key] = auth.value;
        }
    }
    return root;
}

}  // namespace

EngineSelector::EngineSelector(ai_engine::AiEnginePluginManager &aiEnginePluginManager,
                               ModelConfigManager &modelConfigManager)
    : aiEnginePluginManager_(aiEnginePluginManager), modelConfigManager_(modelConfigManager) {}

std::unique_ptr<ai_engine::core_ai::speech::AbstractSpeechEngine> EngineSelector::selectCoreAiSpeechEngine() const {
    auto deployType = selectModelDeployType(AiCapability::Speech);
    const auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Speech, deployType);

    LOGD("modelName {}", modelName);

    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    auto authens = modelConfigManager_.getModelAuthentications(modelName);

    auto engine = aiEnginePluginManager_.createSpeechEngine(modelInfo.model.engineName);
    if (!engine) {
        return nullptr;
    }

    const auto authensJson = authensToEngineConfigJson(authens);

    if (engine->isCloud()) {
        engine->setConfig(authensJson.toStyledString());
    }
    return engine;
}

std::unique_ptr<ai_engine::core_ai::speech::AbstractSpeechEngine> EngineSelector::selectCoreAiSpeechEngine(
    const std::string &modelName) const {
    LOGD("modelName {}", modelName);

    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    auto authens = modelConfigManager_.getModelAuthentications(modelName);

    auto engine = aiEnginePluginManager_.createSpeechEngine(modelInfo.model.engineName);
    if (!engine) {
        return nullptr;
    }

    const auto authensJson = authensToEngineConfigJson(authens);

    if (engine->isCloud()) {
        engine->setConfig(authensJson.toStyledString());
    }
    return engine;
}

std::unique_ptr<ai_engine::core_ai::speech::AbstractSpeechEngine> EngineSelector::selectCoreAiSpeechEngine(
    DeployType deployType) const {
    const auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Speech, deployType);

    LOGD("modelName {}", modelName);

    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    auto authens = modelConfigManager_.getModelAuthentications(modelName);

    auto engine = aiEnginePluginManager_.createSpeechEngine(modelInfo.model.engineName);
    if (!engine) {
        return nullptr;
    }

    const auto authensJson = authensToEngineConfigJson(authens);

    if (engine->isCloud()) {
        engine->setConfig(authensJson.toStyledString());
    }
    return engine;
}

std::unique_ptr<ai_engine::core_ai::vision::AbstractImageProcessEngine> EngineSelector::selectCoreAiImageProcessEngine()
    const {
    return aiEnginePluginManager_.createImageProcessEngine();
}

bool EngineSelector::isLocalCoreAiImageProcessEngineExist() const {
    return aiEnginePluginManager_.createImageProcessEngine() != nullptr;
}

std::unique_ptr<ai_engine::core_ai::vision::AbstractTextRecognitionEngine> EngineSelector::selectTextRecognitionEngine()
    const {
    return aiEnginePluginManager_.createTextRecognitionEngine();
}

std::unique_ptr<ai_engine::core_ai::embedding::AbstractEmbeddingEngine> EngineSelector::selectCoreAiEmbeddingEngine(
    const std::string &modelName) const {
    return aiEnginePluginManager_.createEmbeddingEngine(modelName);
}

std::unique_ptr<ai_engine::gen_ai::nlp::AbstractTextGenerationEngine> EngineSelector::selectGenAiTextGenerationEngine()
    const {
    auto deployType = selectModelDeployType(AiCapability::Nlp);
    LOGI("Get current deploy type: {}.", int(deployType));
    if (deployType == DeployType::Unknown) {
        LOGE("Select nlp engine failed, deploy is unknown.");
        return nullptr;
    }
    auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Nlp, deployType);

    LOGI("Get current model name: {}.", modelName);

    if (modelName.empty()) {
        LOGE("Select nlp engine failed，get current model name is empty.");
        return nullptr;
    }

    return selectGenAiTextGenerationEngine(modelName);
}

std::unique_ptr<ai_engine::gen_ai::nlp::AbstractTextGenerationEngine> EngineSelector::selectGenAiTextGenerationEngine(
    const std::string &modelName) const {
    LOGD("Model name: {}.", modelName);

    if (modelName.empty()) {
        LOGE("Select nlp engine failed, modelName is empty.");
        return nullptr;
    }

    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    if (modelInfo.model.modelName.empty()) {
        LOGE("Select nlp engine failed, modelInfo.model.modelName is empty.");
        return nullptr;
    }

    auto authens = modelConfigManager_.getModelAuthentications(modelName);
    auto engine = aiEnginePluginManager_.createTextGenerationEngine(modelInfo.model.engineName);
    if (engine == nullptr) {
        LOGE("Select nlp engine failed.");
        return nullptr;
    }

    auto authensJson = authensToEngineConfigJson(authens);

    if (engine->isCustomModelEngine()) {
        authensJson["apiUrl"] = modelInfo.model.apiUrl;
        authensJson["modelName"] = modelInfo.model.modelName;
        authensJson["modelVersion"] = modelInfo.model.modelVersion;
    }
    engine->setConfig(authensJson.toStyledString());
    engine->setCurrentModel(modelInfo.model.modelVersion);

    return engine;
}

std::unique_ptr<ai_engine::gen_ai::nlp::AbstractTextGenerationEngine> EngineSelector::selectGenAiTextGenerationEngine(
    DeployType deployType) const {
    if (!modelConfigManager_.getModelDeployEnabled(deployType)) {
        return nullptr;
    }
    auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Nlp, deployType);
    LOGD("Model name: {}.", modelName);

    if (modelName.empty()) {
        LOGE("Select nlp engine failed，get current model name is empty.");
        return nullptr;
    }

    return selectGenAiTextGenerationEngine(modelName);
}

bool EngineSelector::isLocalGenAiTextGenerationEngineExist() const {
    auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Nlp, DeployType::Local);
    if (modelName.empty()) {
        return false;
    }
    return selectGenAiTextGenerationEngine(modelName) != nullptr;
}

std::unique_ptr<ai_engine::gen_ai::vision::AbstractImageGenerationEngine>
EngineSelector::selectGenAiImageGenerationEngine() const {
    auto deployType = selectModelDeployType(AiCapability::Vision);
    if (deployType == DeployType::Unknown) {
        return nullptr;
    }
    const auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Vision, deployType);

    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    auto authens = modelConfigManager_.getModelAuthentications(modelName);
    auto engine = aiEnginePluginManager_.createImageGenerationEngine(modelInfo.model.engineName);
    if (!engine) {
        return nullptr;
    }

    const auto authensJson = authensToEngineConfigJson(authens);

    if (engine->isCloud()) {
        engine->setConfig(authensJson.toStyledString());
    }

    return engine;
}

std::unique_ptr<ai_engine::gen_ai::vision::AbstractImageGenerationEngine>
EngineSelector::selectGenAiImageGenerationEngine(const std::string &modelName) const {
    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    auto engine = aiEnginePluginManager_.createImageGenerationEngine(modelInfo.model.engineName);
    return engine;
}

std::unique_ptr<ai_engine::gen_ai::vision::AbstractImageGenerationEngine>
EngineSelector::selectGenAiImageGenerationEngine(DeployType deployType) const {
    const auto modelName = modelConfigManager_.getCurrentModelName(AiCapability::Vision, deployType);

    ModelInfo modelInfo = modelConfigManager_.getModelInfo(modelName);
    auto engine = aiEnginePluginManager_.createImageGenerationEngine(modelInfo.model.engineName);
    return engine;
}

void EngineSelector::updateCallback(AiCapability capability, DeployType deployType, const std::string &uuid,
                                    const ModelChangedCallbacks::ModelChangedCallback &callback) const {
    return modelConfigManager_.updateCallback(capability, deployType, uuid, callback);
}

void EngineSelector::updateAllDeployTypeCallback(AiCapability capability, const std::string &uuid,
                                                 const ModelChangedCallbacks::ModelChangedCallback &callback) const {
    updateCallback(capability, DeployType::Local, uuid, callback);
    updateCallback(capability, DeployType::PublicCloud, uuid, callback);
    updateCallback(capability, DeployType::Custom, uuid, callback);
}

std::map<DeployType, int> EngineSelector::getCurrentModelSchedulingPriority(AiCapability capability) const {
    int cloudPriority = modelConfigManager_.getModelDeployEnabled(DeployType::PublicCloud) &&
                                !modelConfigManager_.getCurrentModelName(capability, DeployType::PublicCloud).empty()
                            ? modelConfigManager_.getModelSchedulingPriority(DeployType::PublicCloud)
                            : -1;
    int localPriority = modelConfigManager_.getModelDeployEnabled(DeployType::Local) &&
                                !modelConfigManager_.getCurrentModelName(capability, DeployType::Local).empty()
                            ? modelConfigManager_.getModelSchedulingPriority(DeployType::Local)
                            : -1;
    int privatePriority = modelConfigManager_.getModelDeployEnabled(DeployType::Custom) &&
                                  !modelConfigManager_.getCurrentModelName(capability, DeployType::Custom).empty()
                              ? modelConfigManager_.getModelSchedulingPriority(DeployType::Custom)
                              : -1;

    bool publicCloudEnabled = modelConfigManager_.getModelDeployEnabled(DeployType::PublicCloud);
    bool privateCloudEnabled = modelConfigManager_.getModelDeployEnabled(DeployType::Custom);
    bool onDeviceEnabled = modelConfigManager_.getModelDeployEnabled(DeployType::Local);

    if (!publicCloudEnabled) {
        cloudPriority = -1;
    }

    if (!privateCloudEnabled) {
        privatePriority = -1;
    }

    if (!onDeviceEnabled) {
        localPriority = -1;
    }

    std::map<DeployType, int> priorityMap = {
        {DeployType::Local, localPriority},
        {DeployType::PublicCloud, cloudPriority},
        {DeployType::Custom, privatePriority},
    };
    return priorityMap;
}
DeployType EngineSelector::selectModelDeployType(AiCapability capability) const {
    auto priorityMap = getCurrentModelSchedulingPriority(capability);

    int32_t currentPriority{INT32_MAX};
    DeployType currentDeployType{DeployType::Unknown};
    for (const auto &[deployType, priority] : priorityMap) {
        if (priority == -1) {
            continue;
        }
        if (currentPriority > priority) {
            // 当前优先级大于当前[key, value]的优先级，
            // 则更新当前优先级和当前部署类型
            currentPriority = priority;
            currentDeployType = deployType;
        }
    }
    return currentDeployType;
}

DeployType EngineSelector::selectNextPriorityDeployType(AiCapability capability, DeployType currentDeployType) const {
    auto priorityMap = getCurrentModelSchedulingPriority(capability);

    // 获取当前部署类型的优先级
    int currentPriorityValue = priorityMap[currentDeployType];
    int32_t nextPriority{INT32_MAX};
    DeployType nextDeployType{DeployType::Unknown};

    // 找到比当前优先级高（数值更大）的最近优先级
    for (const auto &[deployType, priority] : priorityMap) {
        if (priority == -1 || deployType == currentDeployType) {
            continue;
        }
        if (priority > currentPriorityValue && priority < nextPriority) {
            nextPriority = priority;
            nextDeployType = deployType;
        }
    }

    LOGI("Found next deploy type: {} with priority: {}", static_cast<int>(nextDeployType), nextPriority);
    return nextDeployType;
}

std::unique_ptr<ai_engine::gen_ai::nlp::AbstractTextGenerationEngine>
EngineSelector::selectNextPriorityGenAiTextGenerationEngine(DeployType currentDeployType,
                                                            GenAiTextGenerationEnginePredicate predicate) const {
    auto nextDeployType = selectNextPriorityDeployType(AiCapability::Nlp, currentDeployType);
    if (nextDeployType == DeployType::Unknown) {
        LOGE(
            "Select next priority nlp engine failed, no suitable next deploy "
            "type found.");
        return nullptr;
    }

    LOGI("Selecting next priority NLP engine with deploy type: {}", static_cast<int>(nextDeployType));
    auto engine = selectGenAiTextGenerationEngine(nextDeployType);
    if (predicate(engine)) {
        return engine;
    }

    return selectNextPriorityGenAiTextGenerationEngine(nextDeployType, predicate);
}
