/*
 * 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 "ondevicenlpengine.h"
#include "nlp/llm.h"
#include "utils/logger.h"

#include <jsoncpp/json/json.h>
#include <nlohmann/json.hpp>
#ifndef TRITON_CLIENT
#include <iostream>
#endif

namespace ai_engine::gen_ai::nlp {
OndeviceNlpEngine::OndeviceNlpEngine() : llm_(new LLM()) { loggerNlp::init(); }

OndeviceNlpEngine::~OndeviceNlpEngine() {
    result::AiEngineBaseResult error;
    error = destroy();
};

void OndeviceNlpEngine::setTextGenerationResultCallback(
    nlp::TextGenerationResultCallback callback) {
    chatResultCallback_ = std::move(callback);

    if (llm_ == nullptr) {
        return;
    }

    auto chatCallback = [this](const std::string &result, bool over) {
        this->onChatResult(result, over);
    };
    llm_->setCompletionResultCallback(chatCallback);
}

result::AiEngineBaseResult OndeviceNlpEngine::chat(const std::string &message) {
    result::AiEngineBaseResult error;
    error.success = true;
    chatStopped_ = false;

    if (llm_ == nullptr) {
        error.success = false;
        error.error = ai_engine::EngineError(
            (int)NlpEngineErrorCode::FailedToConnectServer,
            "[本地模型引擎]: NLP对象无效");
        return error;
    }
    if (!inited_) {
        error = init();
        if (!error.success) {
            return error;
        }
        setTextGenerationResultCallback(chatResultCallback_);
    }
    Json::Reader reader;
    Json::Value bodys;
    if (!reader.parse(message, bodys)) {
        error.success = false;
        error.error = ai_engine::EngineError(
            (int)NlpEngineErrorCode::FailedToConnectServer,
            "[本地模型引擎]: NLP消息解析失败");
        return error;
    }
    std::string qus = "\n<|im_start|>";
    for (auto body : bodys["messages"]) {
        qus += body["role"].asString();
        qus += "\n";
        qus += body["content"].asString();
        qus += "<|im_end|>\n<|im_start|>";
    }
    qus += "assistant\n";
    int task_id = -1;
    if (!bodys["task_id"].isNull()) {
        task_id = bodys["task_id"].asInt();
        std::cout << "bodys task_id not null" << std::endl;
    }
    std::cout << qus << std::endl;
    std::cout << "task_id:" << task_id << std::endl;
    struct ChatConfig config;
    if (bodys.isMember("top_k")) {
        config.top_k_input_data = bodys["top_k"].asInt();
    }
    if (bodys.isMember("stream")) {
        config.stream_output = bodys["stream"].asBool();
    }
    llm_->chat(qus, task_id, config);

    return error;
}

void OndeviceNlpEngine::stopChat() {
    if (llm_ == nullptr) {
        return;
    }
    chatStopped_ = true;
    TextGenerationResult endResult;
    endResult.success = true;
    endResult.error = ai_engine::EngineError();
    endResult.modelName = llm_->modelDisplayName_;
    endResult.text = std::string{"{\"result\":\"<|stopChat_8167431|>\",\"is_"
                                 "end\":true,\"finish_reason\":\"stop\"}"};

    llm_->stopChat();
    if (chatResultCallback_ != nullptr) {
        chatResultCallback_(endResult);
    }
}

void OndeviceNlpEngine::onChatResult(const std::string &result, bool over) {
    if (chatResultCallback_ == nullptr) {
        return;
    }

    if (chatStopped_) {
        return;
    }

    Json::Value body;
    std::string finish_reason = "normal";
    body["is_end"] = over;
    body["result"] = result;
    body["finish_reason"] = finish_reason;

    TextGenerationResult res;
    res.success = true;
    res.error = ai_engine::EngineError();
    res.modelName = llm_->modelDisplayName_;
    res.text = body.toStyledString();
    chatResultCallback_(res);
}

result::AiEngineBaseResult OndeviceNlpEngine::init() {
    result::AiEngineBaseResult error;
    llm_->doSetLlmModelName();

    error.success = true;
    if (inited_) {
        error.success = false;
        error.error = ai_engine::EngineError(
            (int)NlpEngineErrorCode::FilePermissionDenied,
            "[本地模型引擎]: NLP模型已经加载");
        return error;
    }
    bool ret = llm_->llm_load_model();
    if (!ret) {
        error.success = false;
        error.error = ai_engine::EngineError(
            (int)NlpEngineErrorCode::LocalNlpModelNotLoaded,
            "[本地模型引擎]: NLP模型加载失败");
        return error;
    }

    inited_ = true;
    return error;
}

result::AiEngineBaseResult OndeviceNlpEngine::destroy() {
    result::AiEngineBaseResult error;
    error.success = true;
    if (!inited_) {
        return error;
    }
    inited_ = false;
    releaseTimerCondition_.notify_all();

    delete llm_;
    llm_ = nullptr;
    return error;
}

namespace nl = nlohmann;

bool OndeviceNlpEngine::setCurrentModel(const std::string &model) {
    bool res = false;
    if (model.size() > 0 && llm_->readySetLlmModelName(model)) {
        res = true;
    } else {
        res = false;
    }

    return res;
}

std::string OndeviceNlpEngine::getLocalModelList() const {
    const char *info = R"(
    {
        "vendor": "",
        "models": [
        ]
    }
    )";
    nl::json llm_json;
    std::string main_info(info);
    nl::json main_json = nl::json::parse(main_info);
    std::string model_info_llm = llm_->GetConfigValue("model_info_list", "llm");

    LOGI("model_info_llm:{}", model_info_llm);
    if (model_info_llm.size() > 0)
        llm_json = nl::json::parse(model_info_llm);
    else
        llm_json = nl::json::parse(main_info);

    main_json["models"].insert(main_json["models"].end(),
                               llm_json["models"].begin(),
                               llm_json["models"].end());
    return main_json.dump(4);
}

std::string OndeviceNlpEngine::currentModel() const {
    if (llm_->targetModelName_ != llm_->modelName_) {
        return llm_->targetModelName_;
    }
    return llm_->modelName_;
}

std::string OndeviceNlpEngine::currentModelName() const {
    return llm_->modelDisplayName_;
}

std::string OndeviceNlpEngine::modelInfo() const {
    if (llm_->clientPreparationStatus_ == ClientPreparationStatus::Success) {
        LOGI("{} prepareModelInfo() clientPreparationStatus_ had Success, "
             "don't try again",
             __LINE__);
    } else if (llm_->clientPreparationStatus_ ==
               ClientPreparationStatus::Failed) {
        LOGE("{} prepareModelInfo() clientPreparationStatus_ had Failed, don't "
             "try again",
             __LINE__);
    } else {
        llm_->prepareModelInfo();
    }
    LOGI("modelInfo:{}", llm_->modelInfo_);
    return llm_->modelInfo_;
}

} // namespace ai_engine::gen_ai::nlp
