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

#include <jsoncpp/json/json.h>
#include <kylin-ai/ai-engine/imageprocessengine.h>

#include "encryption/encryption.h"
#include "utils/logger.h"
#include "utils/util.h"

namespace {
const char *publicCloudSchemaId = "com.kylin.aisdk.model.cloud";
const char *customCloudSchemaId = "com.kylin.aisdk.model.custom";
const char *ondeviceSchemaId = "com.kylin.aisdk.model.ondevice";

const std::string oldModelConfigFile = std::string(getenv("HOME")) + "/.config/model_manager.json";

const std::string baiduNlpModelName = "百度-ERNIE-Bot-4";
const std::string baiduVisionModelName = "百度-文心AI作画高级版";

ModelInfo createModelInfo(const Json::Value &json, AiCapability capability, DeployType deployType,
                          const std::string &engineName, const std::string &vendor) {
    if (!json.isMember("name")) {
        LOGE("Invalid configs format: {}", json.toStyledString());
        return {};
    }
    const std::string modelName = json["name"].asString();
    const std::string model = json["model"].asString();

    if (!json.isMember("subConfig") || !json["subConfig"].isArray()) {
        LOGE("Invalid configs format: {}", json.toStyledString());
        return {};
    }

    std::vector<MultiAuthentication> multiAuthentications;
    for (const auto &config : json["subConfig"]) {
        if (!config.isMember("name") || !config.isMember("config") || !config["config"].isArray()) {
            LOGE("Invalid configs format: {}", json.toStyledString());
            return {};
        }

        const std::string itemName = config["name"].asString();
        const auto configs = config["config"];
        std::vector<Authentication> authentications;
        for (const auto &config : configs) {
            if (!config.isMember("friendlyName") || !config["configName"].isString()) {
                LOGE("Invalid configs format: {}", json.toStyledString());
                return {};
            }
            authentications.push_back({config["friendlyName"].asString(), config["configName"].asString(), ""});
        }

        multiAuthentications.push_back({itemName, authentications});
    }

    const std::string debName = json["debName"].asString();

    const bool is3rdModel = (json["is3rdModel"].asString() == "true");

    const int maxTokens = json["maxTokens"].asInt();

    return {{modelName, model, "", vendor, debName, is3rdModel, capability, deployType, engineName},
            multiAuthentications};
}
std::vector<ModelInfo> createModelInfosFromJsonString(const std::string &jsonString, AiCapability capability,
                                                      DeployType deployType, const std::string &engineName) {
    std::vector<ModelInfo> modelInfos;

    auto root = util::stringToJson(jsonString);
    if (!root.isObject()) {
        LOGE("Invalid configs format: {}", jsonString);
        return {};
    }

    if (!root.isMember("vendor") || !root.isMember("models") || !root["models"].isArray()) {
        LOGE("Invalid configs format: {}", jsonString);
        return {};
    }

    const std::string vendor = root["vendor"].asString();
    const auto models = root["models"];

    for (const auto &model : models) {
        modelInfos.push_back(createModelInfo(model, capability, deployType, engineName, vendor));
    }

    return modelInfos;
}

std::vector<ModelInfo> createCustomModelInfosFromJsonString(const std::string &jsonString,
                                                            const std::string &engineName) {
    std::vector<ModelInfo> modelInfos;

    auto root = util::stringToJson(jsonString);
    if (!root.isObject()) {
        LOGE("Invalid configs format: {}", jsonString);
        return {};
    }

    const Json::Value::Members modelNames = root.getMemberNames();
    for (const auto &modelName : modelNames) {
        if (modelName == "currentModel") {
            continue;
        }

        const auto modelConfig = root[modelName];
        if (!modelConfig.isObject()) {
            LOGE("Invalid configs format for: {}", modelName);
            continue;
        }

        if (!modelConfig.isMember("apiUrl")) {
            LOGE("Invalid configs format: {}", modelConfig.toStyledString());
            continue;
        }

        Model model{modelName,
                    modelConfig["model"].asString(),
                    modelConfig["apiUrl"].asString(),
                    "",
                    "",
                    false,
                    AiCapability::Nlp,
                    DeployType::Custom,
                    engineName};
        std::vector<MultiAuthentication> multiAuthentications{
            {"", {{"APIKey", "apiKey", modelConfig["apiKey"].asString()}}},
        };

        modelInfos.push_back({model, multiAuthentications});
    }

    return modelInfos;
}

Json::Value authenticationsToJsonWithEncryption(const std::vector<MultiAuthentication> &authentications) {
    Json::Value root;
    for (const auto &auth : authentications) {
        Json::Value item;
        item["itemName"] = auth.itemName;
        for (const auto &auth : auth.authentications) {
            Json::Value authenDetails;
            authenDetails["friendlyKeyName"] = auth.friendlyKeyName;
            authenDetails["key"] = auth.key;
            authenDetails["value"] = Encryption::getInstance().encrypt(auth.value);
            item["authentications"].append(authenDetails);
        }
        root.append(item);
    }
    return root;
}

std::vector<MultiAuthentication> jsonToAuthenticationsWithDecryption(const Json::Value &root) {
    std::vector<MultiAuthentication> authentications;
    for (const auto &item : root) {
        MultiAuthentication multiAuth;
        multiAuth.itemName = item["itemName"].asString();
        const auto authens = item["authentications"];
        for (const auto &authDetails : authens) {
            Authentication auth;
            auth.friendlyKeyName = authDetails["friendlyKeyName"].asString();
            auth.key = authDetails["key"].asString();
            auth.value = Encryption::getInstance().decrypt(authDetails["value"].asString());
            multiAuth.authentications.push_back(auth);
        }
        authentications.push_back(multiAuth);
    }
    return authentications;
}

std::vector<MultiAuthentication> customModelJsonToAuthenticationsWithDecryption(const Json::Value &root) {
    if (!root.isMember("apiKey") || !root["apiKey"].isString()) {
        return {};
    }

    const std::string apiKey = root["apiKey"].asString();
    return {{"", {{"APIKey", "apiKey", Encryption::getInstance().decrypt(apiKey)}}}};
}

std::string aiCapabilityString(AiCapability capability) {
    switch (capability) {
        case AiCapability::Nlp:
            return "nlp";
        case AiCapability::Speech:
            return "speech";
        case AiCapability::Vision:
            return "vision";
        default:
            return "";
    }
}

}  // namespace

ModelConfigManager::ModelConfigManager(ai_engine::AiEnginePluginManager &aiEnginePluginManager)
    : aiEnginePluginManager_(aiEnginePluginManager),
      cloudSettings_(publicCloudSchemaId),
      customSettings_(customCloudSchemaId),
      ondeviceSettings_(ondeviceSchemaId) {
    loadModelInfos();
    if (needToMigrateOldModelConfig(baiduNlpModelName)) {
        migrateOldModelConfig(baiduNlpModelName, AiCapability::Nlp);
    }
    if (needToMigrateOldModelConfig(baiduVisionModelName)) {
        migrateOldModelConfig(baiduVisionModelName, AiCapability::Vision);
    }
}

std::vector<Model> ModelConfigManager::getModels(AiCapability capability) const {
    std::vector<Model> models;

    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (const auto &modelInfo : modelInfos_) {
        if (modelInfo.model.capability == capability) {
            // 暂时不显示端侧视觉模型
            if (modelInfo.model.capability == AiCapability::Vision && modelInfo.model.deployType == DeployType::Local) {
                continue;
            }
            models.push_back(modelInfo.model);
        }
    }
    return models;
}

std::vector<Model> ModelConfigManager::getModels(DeployType deployType) const {
    std::vector<Model> models;
    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (const auto &modelInfo : modelInfos_) {
        if (modelInfo.model.deployType == deployType) {
            models.push_back(modelInfo.model);
        }
    }
    return models;
}

std::string ModelConfigManager::getModelVersionByModelName(const std::string &modelName) const {
    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (const auto &modelInfo : modelInfos_) {
        if (modelInfo.model.modelName == modelName) {
            return modelInfo.model.modelVersion;
        }
    }
    return {};
}

std::string ModelConfigManager::getCurrentModelName(AiCapability capability, DeployType deployType) const {
    auto doGetCurrentModelName = [&](const Settings &settings, AiCapability capability) {
        const std::string jsonString = settings.getJsonValue(aiCapabilityString(capability));
        const Json::Value jsonConfig = util::stringToJson(jsonString);
        return jsonConfig.isMember("currentModel") ? jsonConfig["currentModel"].asString() : "";
    };

    switch (deployType) {
        case DeployType::Local: {
            return doGetCurrentModelName(ondeviceSettings_, capability);
        }

        case DeployType::Custom: {
            return doGetCurrentModelName(customSettings_, capability);
        }

        case DeployType::PublicCloud: {
            return doGetCurrentModelName(cloudSettings_, capability);
        }
        default: {
            return "";
        }
    }
    return "";
}

std::vector<MultiAuthentication> ModelConfigManager::getModelAuthentications(const std::string &modelName) const {
    if (!containsModel(modelName)) {
        return {};
    }

    auto modelInfo = getModelInfo(modelName);

    if (modelInfo.model.deployType == DeployType::Local) {
        return {};
    }

    const std::string capabilityKey = aiCapabilityString(modelInfo.model.capability);

    auto &settings = modelInfo.model.deployType == DeployType::Custom ? customSettings_ : cloudSettings_;

    const std::string capabilityJsonString = settings.getJsonValue(capabilityKey);
    Json::Value jsonConfig = util::stringToJson(capabilityJsonString);
    if (!jsonConfig.isMember(modelName)) {
        return modelInfo.multiAuthentications;
    }

    Json::Value jsonAuths = jsonConfig[modelName];

    return modelInfo.model.deployType == DeployType::Custom ? customModelJsonToAuthenticationsWithDecryption(jsonAuths)
                                                            : jsonToAuthenticationsWithDecryption(jsonAuths);
}

bool ModelConfigManager::getModelDeployEnabled(DeployType deployType) const {
    switch (deployType) {
        case DeployType::Local: {
            return ondeviceSettings_.getBoolValue("enabled");
        }
        case DeployType::Custom: {
            return customSettings_.getBoolValue("enabled");
        }
        case DeployType::PublicCloud: {
            return cloudSettings_.getBoolValue("enabled");
        }
        default: {
            return false;
        }
    }
    return false;
}

bool ModelConfigManager::setCurrentModelName(AiCapability capability, DeployType deployType,
                                             const std::string &modelName) {
    if (!containsModel(modelName)) {
        LOGE("Model not found: {}", modelName);
        return false;
    }

    modelChangedCallbacks_.callAllCallbacks(capability, deployType);

    auto doSetCurrentModelName = [modelName](Settings &settings, AiCapability capability) {
        const std::string jsonString = settings.getJsonValue(aiCapabilityString(capability));
        Json::Value jsonConfig = util::stringToJson(jsonString);
        jsonConfig["currentModel"] = modelName;
        const std::string newJsonString = jsonConfig.toStyledString();
        settings.setJsonValue(aiCapabilityString(capability), newJsonString);
    };

    switch (deployType) {
        case DeployType::Local: {
            doSetCurrentModelName(ondeviceSettings_, capability);
            return true;
        }
        case DeployType::Custom: {
            doSetCurrentModelName(customSettings_, capability);
            return true;
        }
        case DeployType::PublicCloud: {
            doSetCurrentModelName(cloudSettings_, capability);
            return true;
        }
        default:
            return false;
    }
    return false;
}

/// !todo 按照authentications中指定的itemName进行更新
bool ModelConfigManager::setModelAuthentications(const std::string &modelName,
                                                 const std::vector<MultiAuthentication> &authentications) {
    if (authentications.empty()) {
        LOGE("Authentications is empty.");
        return false;
    }

    if (!containsModel(modelName)) {
        LOGE("Model not found: {}.", modelName);
        return false;
    }

    auto modelInfo = getModelInfo(modelName);
    if (modelInfo.model.deployType == DeployType::Local) {
        return false;
    }

    auto &settings = modelInfo.model.deployType == DeployType::Custom ? customSettings_ : cloudSettings_;

    updateModelAuthenticationsSettings(settings, modelInfo, authentications);
    updateModelAuthenticationsInfo(modelName, authentications);

    modelChangedCallbacks_.callAllCallbacks(modelInfo.model.capability, modelInfo.model.deployType);

    return true;
}

void ModelConfigManager::updateModelAuthenticationsSettings(Settings &settings, const ModelInfo &modelInfo,
                                                            const std::vector<MultiAuthentication> &authentications) {
    const std::string capabilityKey = aiCapabilityString(modelInfo.model.capability);

    // 对于custom类型返回的是字符串，其他类型返回Json::Value
    auto jsonAuths = modelInfo.model.deployType == DeployType::Custom
                         ? Encryption::getInstance().encrypt(
                               authentications.empty() ? "" : authentications[0].authentications[0].value)
                         : authenticationsToJsonWithEncryption(authentications);

    const std::string capabilityJsonString = settings.getJsonValue(capabilityKey);
    Json::Value jsonConfig = util::stringToJson(capabilityJsonString);

    auto &keyObj = modelInfo.model.deployType == DeployType::Custom ? jsonConfig[modelInfo.model.modelName]["apiKey"]
                                                                    : jsonConfig[modelInfo.model.modelName];
    LOGD("Update model authentications settings: {}.", keyObj.toStyledString());
    keyObj = jsonAuths;
    const std::string jsonString = jsonConfig.toStyledString();
    LOGD("Update model authentications settings: {}.", jsonConfig.toStyledString());
    settings.setJsonValue(capabilityKey, jsonString);
}

void ModelConfigManager::updateModelAuthenticationsInfo(const std::string &modelName,
                                                        const std::vector<MultiAuthentication> &authentications) {
    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (auto &model : modelInfos_) {
        LOGD("{} {}", model.model.modelName, modelName);
        if (model.model.modelName == modelName) {
            model.multiAuthentications = authentications;
            break;
        }
    }
}

void ModelConfigManager::clearModelAuthenticationsSettings(Settings &settings, const ModelInfo &modelInfo) {
    const std::string capabilityKey = aiCapabilityString(modelInfo.model.capability);
    const std::string capabilityJsonString = settings.getJsonValue(capabilityKey);
    LOGD("{}", capabilityJsonString);
    Json::Value jsonConfig = util::stringToJson(capabilityJsonString);

    if (modelInfo.model.deployType == DeployType::Custom) {
        jsonConfig[modelInfo.model.modelName]["apiKey"] = "";
        settings.setJsonValue(capabilityKey, jsonConfig.toStyledString());
        return;
    }

    if (jsonConfig.isMember(modelInfo.model.modelName)) {
        auto &authentications = jsonConfig[modelInfo.model.modelName];
        if (authentications.isArray()) {
            for (auto &authentication : authentications) {
                for (auto &auth : authentication["authentications"]) {
                    auth["value"] = "";
                }
            }
        }
        if (jsonConfig["currentModel"].asString() == modelInfo.model.modelName) {
            jsonConfig.removeMember("currentModel");
        }
        std::string jsonString = jsonConfig.toStyledString();
        LOGD("jsonString {}", jsonString);
        settings.setJsonValue(capabilityKey, jsonString);
    }
}

void ModelConfigManager::clearModelAuthenticationsInfo(const std::string &modelName) {
    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (auto &model : modelInfos_) {
        if (model.model.modelName == modelName) {
            model.multiAuthentications.clear();
            break;
        }
    }
}

// 对于云端模型来说是删除所有的模型配置
// 对于自定义模型仅仅是清除apiKey的内容
/// !todo 具体逻辑需要进一步优化
bool ModelConfigManager::clearModelAuthentications(const std::string &modelName) {
    if (!containsModel(modelName)) {
        LOGE("Model not found: {}", modelName);
        return false;
    }

    ModelInfo modelInfo = getModelInfo(modelName);
    auto &settings = modelInfo.model.deployType == DeployType::Custom ? customSettings_ : cloudSettings_;

    clearModelAuthenticationsSettings(settings, modelInfo);
    clearModelAuthenticationsInfo(modelName);

    modelChangedCallbacks_.callAllCallbacks(modelInfo.model.capability, modelInfo.model.deployType);

    return true;
}

bool ModelConfigManager::setModelSchedulingPriority(DeployType deployType, int priority) {
    modelChangedCallbacks_.callAllCallbacks(AiCapability::Nlp, deployType);
    modelChangedCallbacks_.callAllCallbacks(AiCapability::Speech, deployType);
    modelChangedCallbacks_.callAllCallbacks(AiCapability::Vision, deployType);

    switch (deployType) {
        case DeployType::Local: {
            ondeviceSettings_.setIntValue("priority", priority);
            return true;
        }
        case DeployType::Custom: {
            customSettings_.setIntValue("priority", priority);
            return true;
        }
        case DeployType::PublicCloud: {
            cloudSettings_.setIntValue("priority", priority);
            return true;
        }
    }
    return false;
}

int ModelConfigManager::getModelSchedulingPriority(DeployType deployType) const {
    switch (deployType) {
        case DeployType::Local: {
            return ondeviceSettings_.getIntValue("priority");
        }
        case DeployType::Custom: {
            return customSettings_.getIntValue("priority");
        }
        case DeployType::PublicCloud: {
            return cloudSettings_.getIntValue("priority");
        }
    }
    return -1;
}

bool ModelConfigManager::setModelDeployEnabled(DeployType deployType, bool enabled) {
    modelChangedCallbacks_.callAllCallbacks(AiCapability::Nlp, deployType);
    modelChangedCallbacks_.callAllCallbacks(AiCapability::Speech, deployType);
    modelChangedCallbacks_.callAllCallbacks(AiCapability::Vision, deployType);

    switch (deployType) {
        case DeployType::Local:
            return ondeviceSettings_.setBoolValue("enabled", enabled);
        case DeployType::Custom:
            return customSettings_.setBoolValue("enabled", enabled);
        case DeployType::PublicCloud:
            return cloudSettings_.setBoolValue("enabled", enabled);
    }
    return false;
}

bool ModelConfigManager::needToMigrateOldModelConfig(const std::string &modelName) {
    if (util::readFile(oldModelConfigFile).empty()) {
        return false;
    }

    const auto autehns = getModelAuthentications(modelName);
    return autehns.empty() || autehns.back().authentications.back().value.empty();
}

// 旧版本的配置只有百度的，并且只有nlp的vision的能力
void ModelConfigManager::migrateOldModelConfig(const std::string &modelName, AiCapability capability) {
    std::vector<ModelInfo> odlModelInfos;
    const std::string oldVisionModelName = "百度-文心AI作画高级版";
    const std::string modelInfo = util::readFile(oldModelConfigFile);
    const Json::Value root = util::stringToJson(modelInfo);
    // 百度
    if (!root.isMember("0")) {
        return;
    }
    Json::Value baiduModelInfo = root["0"];
    const std::string capabilityKey = capability == AiCapability::Nlp ? "0" : "1";
    if (!baiduModelInfo.isMember(capabilityKey)) {
        return;
    }

    Json::Value nlpModelInfo = baiduModelInfo[capabilityKey];
    const std::string apiKey = nlpModelInfo["apiKey"].asString();
    const std::string secretKey = nlpModelInfo["secretKey"].asString();
    const std::string currentModelName = getCurrentModelName(capability, DeployType::PublicCloud);
    setModelAuthentications(modelName, {{"",
                                         {{"APIKey", "apiKey", Encryption::getInstance().decrypt(apiKey)},
                                          {"SecretKey", "secretKey", Encryption::getInstance().decrypt(secretKey)}}}});
    if (currentModelName.empty()) {
        setCurrentModelName(capability, DeployType::PublicCloud, modelName);
    }
}

void ModelConfigManager::loadModelInfos() {
    auto allEngineNames = aiEnginePluginManager_.allEngineName();
    for (const auto &engineName : allEngineNames) {
        loadNlpModelInfos(engineName);
        loadSpeechModelInfos(engineName);
        loadVisionModelInfos(engineName);
    }

    for (auto &model : modelInfos_) {
        LOGI("Model {} loaded", model.model.modelName);
    }
}

void ModelConfigManager::loadNlpModelInfos(const std::string &engineName) {
    auto engine = aiEnginePluginManager_.createTextGenerationEngine(engineName);
    if (!engine) {
        return;
    }

    std::vector<ModelInfo> models;
    if (engine->isCloud() && !engine->isCustomModelEngine()) {
        auto cloudEngine = static_cast<ai_engine::gen_ai::nlp::AbstractTextGenerationEngine *>(engine.get());
        models = createModelInfosFromJsonString(cloudEngine->modelInfo(), AiCapability::Nlp, DeployType::PublicCloud,
                                                engineName);
    } else if (engine->isCustomModelEngine()) {
        models = createCustomModelInfosFromJsonString(customSettings_.getJsonValue("nlp"), engineName);
    } else if (!engine->isCloud()) {
        models = createModelInfosFromJsonString(engine->modelInfo(), AiCapability::Nlp, DeployType::Local, engineName);
        std::string jsonString = ondeviceSettings_.getJsonValue(aiCapabilityString(AiCapability::Nlp));
        Json::Value config = util::stringToJson(jsonString);
        if (!config.isMember("currentModel")) {
            config["currentModel"] = "";
        }
        // 如果为空，且已安装本地模型，则设置默认值
        if (config["currentModel"].asString().empty()) {
            if (!models.empty()) {
                config["currentModel"] = models[0].model.modelName;
            }
        }
        ondeviceSettings_.setJsonValue(aiCapabilityString(AiCapability::Nlp), util::jsonToString(config));
    }

    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    modelInfos_.insert(modelInfos_.end(), models.begin(), models.end());
}

void ModelConfigManager::loadSpeechModelInfos(const std::string &engineName) {
    auto engine = aiEnginePluginManager_.createSpeechEngine(engineName);
    if (engine && engine->isCloud()) {
        auto cloudEngine = static_cast<ai_engine::core_ai::speech::AbstractSpeechEngine *>(engine.get());
        auto models = createModelInfosFromJsonString(cloudEngine->modelInfo(), AiCapability::Speech,
                                                     DeployType::PublicCloud, engineName);

        std::lock_guard<std::mutex> lock(modelInfosMutex_);
        modelInfos_.insert(modelInfos_.end(), models.begin(), models.end());
    }
}

void ModelConfigManager::loadVisionModelInfos(const std::string &engineName) {
    auto engine = aiEnginePluginManager_.createImageGenerationEngine(engineName);
    if (!engine) {
        return;
    }

    if (engine->isCloud()) {
        auto cloudEngine = static_cast<ai_engine::gen_ai::vision::AbstractImageGenerationEngine *>(engine.get());
        auto models = createModelInfosFromJsonString(cloudEngine->modelInfo(), AiCapability::Vision,
                                                     DeployType::PublicCloud, engineName);

        std::lock_guard<std::mutex> lock(modelInfosMutex_);
        modelInfos_.insert(modelInfos_.end(), models.begin(), models.end());
    } else {
        auto models =
            createModelInfosFromJsonString(engine->modelInfo(), AiCapability::Vision, DeployType::Local, engineName);

        std::lock_guard<std::mutex> lock(modelInfosMutex_);
        modelInfos_.insert(modelInfos_.end(), models.begin(), models.end());
    }
}

bool ModelConfigManager::containsModel(const std::string &modelName) const {
    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    return std::find_if(modelInfos_.begin(), modelInfos_.end(), [&modelName](const ModelInfo &m) {
               return m.model.modelName == modelName;
           }) != modelInfos_.end();
}

ModelInfo ModelConfigManager::getModelInfo(const std::string &modelName) const {
    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    auto it = std::find_if(modelInfos_.begin(), modelInfos_.end(),
                           [&modelName](const ModelInfo &m) { return m.model.modelName == modelName; });

    if (it != modelInfos_.end()) {
        return *it;
    }

    return {};
}

bool ModelConfigManager::addCustomModel(const std::string &modelName, const std::string &model,
                                        const std::string &apiKey, const std::string &apiUrl) {
    if (containsModel(modelName)) {
        LOGE("该模型已被添加 {}", modelName);
        return false;
    }
    Json::Value currentCustomJsonConfig = util::stringToJson(customSettings_.getJsonValue("nlp"));
    currentCustomJsonConfig[modelName]["name"] = modelName;
    currentCustomJsonConfig[modelName]["apiKey"] = Encryption::getInstance().encrypt(apiKey);
    currentCustomJsonConfig[modelName]["apiUrl"] = apiUrl;
    currentCustomJsonConfig[modelName]["model"] = model;
    LOGD("AddCustomModel: {}", currentCustomJsonConfig.toStyledString());
    // 当前自定义模型只支持nlp模型
    customSettings_.setJsonValue("nlp", currentCustomJsonConfig.toStyledString());

    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    modelInfos_.push_back({{modelName, model, apiUrl, "", "", false, AiCapability::Nlp, DeployType::Custom, "custom"},
                           {{"", {{"APIKey", "apiKey", apiKey}}}}});
    return true;
}

bool ModelConfigManager::deleteCustomModel(const std::string &modelName) {
    Json::Value currentCustomJsonConfig = util::stringToJson(customSettings_.getJsonValue("nlp"));
    if (!currentCustomJsonConfig.isMember(modelName)) {
        return false;
    }
    currentCustomJsonConfig.removeMember(modelName);
    if (currentCustomJsonConfig.isMember("currentModel") && currentCustomJsonConfig["currentModel"].isString() &&
        currentCustomJsonConfig["currentModel"].asString() == modelName) {
        currentCustomJsonConfig["currentModel"] = "";
    }

    customSettings_.setJsonValue("nlp", currentCustomJsonConfig.toStyledString());

    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    auto it = std::remove_if(modelInfos_.begin(), modelInfos_.end(),
                             [&modelName](const ModelInfo &m) { return m.model.modelName == modelName; });

    modelInfos_.erase(it, modelInfos_.end());

    return true;
}

bool ModelConfigManager::setCustomModelApiUrl(const std::string &modelName, const std::string &apiUrl) {
    Json::Value currentCustomJsonConfig = util::stringToJson(customSettings_.getJsonValue("nlp"));
    if (!currentCustomJsonConfig.isMember(modelName)) {
        return false;
    }
    currentCustomJsonConfig[modelName]["apiUrl"] = apiUrl;
    LOGI("SetCustomModelApiUrl: {}", currentCustomJsonConfig.toStyledString());
    customSettings_.setJsonValue("nlp", currentCustomJsonConfig.toStyledString());

    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (auto &model : modelInfos_) {
        if (model.model.modelName == modelName) {
            model.model.apiUrl = apiUrl;
            break;
        }
    }
    return true;
}

bool ModelConfigManager::setCustomModelVersion(const std::string &modelName, const std::string &modelVersion) {
    Json::Value currentCustomJsonConfig = util::stringToJson(customSettings_.getJsonValue("nlp"));
    if (!currentCustomJsonConfig.isMember(modelName)) {
        return false;
    }

    currentCustomJsonConfig[modelName]["model"] = modelVersion;
    customSettings_.setJsonValue("nlp", currentCustomJsonConfig.toStyledString());

    std::lock_guard<std::mutex> lock(modelInfosMutex_);
    for (auto &model : modelInfos_) {
        if (model.model.modelName == modelName) {
            model.model.modelVersion = modelVersion;
            break;
        }
    }
    return true;
}

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