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

#include <sstream>
#include <nlohmann/json.hpp>

#ifdef TRITON_CLIENT
namespace tc = triton::client;
#endif
namespace nl = nlohmann;

#ifdef TRITON_CLIENT
#define FAIL_IF_ERR(X, MSG)                                                    \
    {                                                                          \
        tc::Error err = (X);                                                   \
        if (!err.IsOk()) {                                                     \
            std::cerr << "error: " << (MSG) << ": " << err << std::endl;       \
        }                                                                      \
    }

void LLM::llm_callback(tc::InferResult *result) {
    std::shared_ptr<tc::InferResult> result_ptr(result);
    std::lock_guard<std::mutex> lk(mutex_);
    result_list_.push_back(result_ptr);
    cv_.notify_all();
}
#endif

bool LLM::llm_load_model() {

#ifdef TRITON_CLIENT
    if (clientPreparationStatus_ == ClientPreparationStatus::Success) {
        LOGI("{} prepareModelInfo() clientPreparationStatus_ had Success, "
             "don't try again",
             __LINE__);
    } else if (clientPreparationStatus_ == ClientPreparationStatus::Failed) {
        LOGE("{} prepareModelInfo() clientPreparationStatus_ had Failed, don't "
             "try again",
             __LINE__);
    } else {
        prepareModelInfo();
    }
    bool nlpModelReady = false;
    tc::Error err_modelReady =
        loadClient_->IsModelReady(&nlpModelReady, modelName_);
    if (!err_modelReady.IsOk()) {
        LOGE("{} err_modelReady.Message(): {}", __LINE__,
             err_modelReady.Message());
        return false;
    }
    LOGI("{} nlpModelReady:{}", __LINE__, nlpModelReady);
    if (!nlpModelReady) {
        tc::Error err_loadmodel =
            loadClient_->LoadModel(modelName_, {}, {}, {}, modelLoadTimeoutUs_);
        if (!err_loadmodel.IsOk()) {
            LOGE("{} err_loadmodel.Message(): {}", __LINE__,
                 err_loadmodel.Message());
            return false;
        }
    }

    modelInfoUpdate();

    nl::json modelInfoJson = nl::json::parse(modelInfo_);
    try {
        for (auto &it_json : modelInfoJson["models"]) {
            if (it_json["model"] == modelName_) {
                modelDisplayName_ = it_json["name"].get<std::string>();
                break;
            }
        }
    } catch (const nl::json::parse_error &e) {
        // 捕获解析错误
        LOGE("Parse error at byte {}: {}", e.byte, e.what());
        modelDisplayName_ = "";
    } catch (const nl::json::exception &e) {
        // 捕获所有来自nlohmann::json库的异常
        LOGE("JSON exception: {}", e.what());
        modelDisplayName_ = "";
    } catch (const std::exception &e) {
        // 捕获所有标准C++异常
        LOGE("Standard exception: {}", e.what());
        modelDisplayName_ = "";
    }
#endif
    return true;
}

LLM::LLM() {
#ifdef TRITON_CLIENT
    bool verbose = false;
    modelLoadTimeoutUs_ = 60 * 1000 * 1000;
#if 1
    tc::Error err_create = tc::InferenceServerGrpcClient::Create(
        &loadClient_, grpcServerUrl_, verbose);
    if (!err_create.IsOk()) {
        LOGE("{} err_create.Message(): {}", __LINE__, err_create.Message());
    }

    tc::Error err_create_grpc = tc::InferenceServerGrpcClient::Create(
        &inferClient_, grpcServerUrl_, verbose);
    if (!err_create_grpc.IsOk()) {
        LOGE("{} err_create_grpc.Message(): {}", __LINE__,
             err_create_grpc.Message());
    }
#else
    bool use_ssl = false;
    bool use_cached_channel = true;
    tc::SslOptions ssl_options = tc::SslOptions();
    tc::Error err_create = KyInfer::Client::CreateGrpcClient(
        &loadClient_, grpcServerUrl_, verbose, use_ssl, ssl_options,
        tc::KeepAliveOptions(), use_cached_channel);
#endif
#endif
}

bool LLM::prepareModelInfo() {
#ifdef TRITON_CLIENT
    clientPreparationStatus_ = ClientPreparationStatus::Success;
    // 此函数仅尝试运行一次，若异常则说明环境不可用，不再多次尝试，避免后续影响其他模块
    tc::Error err_loadmodel;
    int reTryTime = 10;
    std::string modelBank = "model_bank";
    if (!modelBankIsReady_) {
        do {
            LOGI("model bank is not ready, we need load");
            err_loadmodel = loadClient_->LoadModel(modelBank);
            if (!err_loadmodel.IsOk()) {
                LOGE("{} LoadModel err_loadmodel.Message():{}", __LINE__,
                     err_loadmodel.Message());
                sleep(1);
                reTryTime--;
            } else {
                modelBankIsReady_ = true;
            }
        } while (!err_loadmodel.IsOk() && reTryTime > 0);
    } else {
        LOGI("model bank is ready, we don't need load");
    }

    modelName_ = GetConfigValue("llm", "llm");
    if (modelName_.size() < 1) {
        LOGE("line:{} model_name.size():{}", __LINE__, modelName_.size());
        modelName_ = "llm_Qwen-2.5-3b_1.0";
    }

    auto now = std::chrono::system_clock::now();
    auto value = now.time_since_epoch();
    auto microseconds =
        std::chrono::duration_cast<std::chrono::microseconds>(value).count();
    double seconds = microseconds / 1000000.0;
    request_id_ = std::to_string(seconds);
    LOGI("{}: request_id_:{}\n", __LINE__, request_id_);
    tc::Headers http_headers;
    uint32_t stream_timeout = 0;
    tc::Error err_start_stream = inferClient_->StartStream(
        std::bind(&LLM::llm_callback, this, std::placeholders::_1),
        false /*ship_stats*/, stream_timeout, http_headers);
    if (!err_start_stream.IsOk()) {
        LOGE("{} err_start_stream.Message():{}", __LINE__,
             err_start_stream.Message());
        // return false;
    }
    default_model_info_ = R"(
        {
            "vendor": "",
            "models": [{
                "model": "llm_Qwen-2.5-3b_1.0",
                "name": "Qwen-2.5-3b_1.0",
                "capability": "llm",
                "subConfig": [],
                "debName": "kylin-qwen2.5-3b-gguf-model",
                "is3rdModel": "false",
                "maxTokens": 1024
            }]
        })";

    modelInfo_ = R"(
      {
          "vendor": "",
          "models": []
      })";

    n_predict_input_ = 1024;
    modelInfoUpdate();
    LOGI("modelInfo_:{}\n", modelInfo_);
#endif
    return true;
}

void LLM::modelInfoUpdate() {
#ifdef TRITON_CLIENT
    std::string model_info_list = GetConfigValue("model_info_list", "llm");

    nl::json j = nl::json::parse(modelInfo_);
    try {
        j["models"].clear();
        nl::json new_models = nl::json::parse(model_info_list)["models"];
        j["models"].insert(j["models"].end(), new_models.begin(),
                           new_models.end());

        for (auto &it_json : j["models"]) {
            // 获取模型配置
            inference::ModelConfigResponse config;
            tc::Error err_modelconfig =
                loadClient_->ModelConfig(&config, it_json["model"]);

            if (!err_modelconfig.IsOk()) {
                LOGE("{} unable to get ModelConfig, err_modelconfig.Message(): "
                     "{}",
                     __LINE__, err_modelconfig.Message());
            } else {
                // 查找名为 'n_ctx' 的参数
                int n_ctx = 0;
                auto it = config.config().parameters().find("n_ctx");
                if (it != config.config().parameters().end()) {
                    LOGI("n_ctx parameter value: {}",
                         it->second.string_value());
                    n_ctx = std::stoi(it->second.string_value());
                } else {
                    LOGI("n_ctx parameter not found.");
                    n_ctx = 0;
                }
                it_json["maxTokens"] = n_ctx;

                auto it_n_predic =
                    config.config().parameters().find("n_predict");
                if (it_n_predic != config.config().parameters().end()) {
                    LOGI("n_predict_input parameter value: {}",
                         it_n_predic->second.string_value());

                    n_predict_input_ =
                        std::stoi(it_n_predic->second.string_value());
                } else {
                    LOGI("n_predict_input parameter not found. set to 1024");
                    n_predict_input_ = 1024;
                }
            }
        }
        modelInfo_ = j.dump(4);
        // 成功解析后的代码逻辑
    } catch (const nl::json::parse_error &e) {
        // 捕获解析错误
        LOGE("Parse error at byte {}: {}", e.byte, e.what());
        modelInfo_ = default_model_info_;
    } catch (const nl::json::exception &e) {
        // 捕获所有来自nlohmann::json库的异常
        LOGE("JSON exception: {}", e.what());
        modelInfo_ = default_model_info_;
    } catch (const std::exception &e) {
        // 捕获所有标准C++异常
        LOGE("Standard exception: {}", e.what());
        modelInfo_ = default_model_info_;
    }

    if (modelInfo_.size() < 1) {
        LOGE("modelInfo_.size() < 1 : {}", modelInfo_.size());
        modelInfo_ = default_model_info_;
    }
#endif
}

LLM::~LLM() {
    // TODO:
}

bool LLM::doSetLlmModelName() {
#ifdef TRITON_CLIENT

    if (clientPreparationStatus_ == ClientPreparationStatus::Success) {
        LOGI("{} prepareModelInfo() clientPreparationStatus_ had Success, "
             "don't try again",
             __LINE__);
    } else if (clientPreparationStatus_ == ClientPreparationStatus::Failed) {
        LOGE("{} prepareModelInfo() clientPreparationStatus_ had Failed, don't "
             "try again",
             __LINE__);
    } else {
        prepareModelInfo();
    }
    nl::json llm_json;
    if (modelInfo_.size() > 0)
        llm_json = nl::json::parse(modelInfo_);
    else if (targetModelName_ == modelName_) {
        return true;
    } else {
        return false;
    }
    for (auto llm_obj : llm_json["models"]) {
        if ((modelName_ != targetModelName_) &&
            (targetModelName_ == llm_obj["model"])) {
            bool ready = false;
            loadClient_->IsModelReady(&ready, modelName_);
            // ready 用于判断当前engine是否已经完成模型加载
            // 如果未加载，说明是在模型加载前设置要加载的模型名，此时仅用于设置模型名，后续流程会加载模型
            // 如果已经加载，则是在运行中切换模型，需要将已加载的模型卸载后再加载
            if (ready) {
                loadClient_->UnloadModel(modelName_);
            }
            modelName_ = targetModelName_;
            if (ready) {
                // 增加加载模型超时功能，modelLoadTimeoutUs_秒未返回则超时
                loadClient_->LoadModel(modelName_, {}, {}, {},
                                       modelLoadTimeoutUs_);
            }
            modelInfoUpdate();
            return true;
        }
    }
#endif
    return false;
}

bool LLM::readySetLlmModelName(const std::string &modelName) {
    targetModelName_ = modelName;
    return true;
}

void LLM::setCompletionResultCallback(CompletionResultCallback callback) {
    completionCallbacks_ = std::move(callback);
}

void LLM::stopChat() {
#ifdef TRITON_CLIENT
    if (inferClient_ == nullptr ||
        !is_chatting_.load(std::memory_order_acquire)) {
        return;
    }
    try {
        struct ChatConfig config;
        // int lora_count = getLoraCount();
        int lora_count =
            2; // 为规避triton接收无lora的推理请求会报错的问题做出的临时修改
        std::vector<float> lora_scale_input_data(lora_count, 0.0f);
        setLoraScale(-1, lora_scale_input_data);
        config.request_type_data = 5;
        config.request_id = request_id_;
        config.top_k_input_data = 0;
        config.top_p_input_data = 0.0;
        config.n_predict_input_data = 0;
        config.lora_scale_input_data = lora_scale_input_data;
        std::vector<tc::InferInput *> inputs;

        chat_inputs(config, inputs);

        return;
    } catch (const std::exception &e) {
        LOGE("Caught exception: {}", e.what());
        // 捕获所有从 std::exception 派生的异常
    }
#endif
}

std::string LLM::GetConfigValue(const char *key, const char *engine) {
#ifdef TRITON_CLIENT
    tc::Error err_loadmodel = loadClient_->LoadModel("model_bank");
    if (!err_loadmodel.IsOk()) {
        LOGE("{} err_loadmodel.Message(): {}", __LINE__,
             err_loadmodel.Message());
        return "";
    }
    std::string model_name = "model_bank";
    tc::InferOptions options(model_name);
    options.model_version_ = "";
    options.priority_ = 1;

    std::vector<std::string> input_key_data{std::string(key)};
    std::vector<std::string> input_engine_data{std::string(engine)};

    tc::InferInput *input0, *input1;
    std::vector<int64_t> shape{1};
    tc::InferInput::Create(&input0, "key", shape, "BYTES");
    tc::InferInput::Create(&input1, "engine", shape, "BYTES");
    std::shared_ptr<tc::InferInput> input0_ptr, input1_ptr;
    input0_ptr.reset(input0);
    input1_ptr.reset(input1);
    input0_ptr->AppendFromString(input_key_data);
    input1_ptr->AppendFromString(input_engine_data);
    std::vector<tc::InferInput *> inputs = {input0_ptr.get(), input1_ptr.get()};

    tc::InferRequestedOutput *output0;
    tc::Error err_out0 = tc::InferRequestedOutput::Create(&output0, "value");
    std::shared_ptr<tc::InferRequestedOutput> output0_ptr;
    output0_ptr.reset(output0);
    std::vector<const tc::InferRequestedOutput *> outputs = {output0_ptr.get()};

    tc::InferResult *results;
    tc::Headers http_headers;
    tc::Error err_infer =
        loadClient_->Infer(&results, options, inputs, outputs, http_headers);
    if (!err_infer.IsOk()) {
        LOGE("{} err_infer.Message(): {}", __LINE__, err_infer.Message());
        return "";
    }

    std::shared_ptr<tc::InferResult> results_ptr;
    results_ptr.reset(results);

    char *output0_data_ptr;
    size_t recv_output0_byte_size;

    std::vector<std::string> text_output_data;
    tc::Error err_paser = results_ptr->StringData("value", &text_output_data);
    if (!err_paser.IsOk()) {
        LOGE("{} err_paser.Message(): {}", __LINE__, err_paser.Message());
        return "";
    }
    return text_output_data[0];
#else
    return "";
#endif
}

int LLM::setLoraScale(int task_id, std::vector<float> &lora_scale_input_data) {
    // 检查task_id是否是有效索引
    if (task_id >= 0 &&
        task_id < static_cast<int>(lora_scale_input_data.size())) {
        // 填充所有元素为0
        std::fill(lora_scale_input_data.begin(), lora_scale_input_data.end(),
                  0.0f);

        // 设置指定索引的值为1
        lora_scale_input_data[task_id] = 1.0f;
    } else {
        return -1;
    }
    return 0;
}

int LLM::getLoraCount() {
#ifdef TRITON_CLIENT
    // 获取模型配置
    inference::ModelConfigResponse config;
    tc::Error err_modelconfig = loadClient_->ModelConfig(&config, modelName_);

    if (!err_modelconfig.IsOk()) {
        LOGE("{} unable to get ModelConfig, err_modelconfig.Message(): {}",
             __LINE__, err_modelconfig.Message());
        return -1;
    }

    // 查找名为 'lora' 的参数
    int semicolon_count = 0;
    auto it = config.config().parameters().find("lora");
    if (it != config.config().parameters().end()) {
        LOGI("LoRA parameter value: {}", it->second.string_value());
        if (it->second.string_value().empty())
            return 0;

        std::istringstream stream(it->second.string_value());
        std::string segment;
        int count = 0;

        while (std::getline(stream, segment, ';')) {
            if (!segment.empty()) { // 忽略空段
                ++count;
            }
        }

        // 如果最后一个字符是分隔符，则原始字符串可能有一个额外的空段
        if (!it->second.string_value().empty() &&
            it->second.string_value().back() == ';') {
            ++count;
        }

        semicolon_count = count;
    } else {
        LOGE("LoRA parameter not found.");
        semicolon_count = 0;
    }

    return semicolon_count;
#else
    return 0;
#endif
}
#ifdef TRITON_CLIENT

template <typename T>
std::shared_ptr<tc::InferInput> createInput( // 返回智能指针
    const std::string &name, const std::vector<int64_t> &shape,
    const std::string &datatype, const T &data) {
    tc::InferInput *raw_input;
    FAIL_IF_ERR(tc::InferInput::Create(&raw_input, name.c_str(), shape,
                                       datatype.c_str()),
                "unable to create '" + name + "'");

    std::shared_ptr<tc::InferInput> input_ptr(raw_input);
    FAIL_IF_ERR(input_ptr->Reset(), "unable to reset '" + name + "'");

    if constexpr (std::is_same_v<T, std::vector<std::string>>) {
        input_ptr->SetBinaryData(true);
        FAIL_IF_ERR(input_ptr->AppendFromString(data),
                    "unable to set data for '" + name + "'");
    } else if constexpr (std::is_same_v<T, std::vector<float>>) {
        FAIL_IF_ERR(
            input_ptr->AppendRaw(reinterpret_cast<const uint8_t *>(data.data()),
                                 sizeof(float) * data.size()),
            "unable to set data for '" + name + "'");
    } else {
        FAIL_IF_ERR(input_ptr->AppendRaw(
                        reinterpret_cast<const uint8_t *>(&data), sizeof(T)),
                    "unable to set data for '" + name + "'");
    }

    LOGI("Created input: {}", input_ptr->Name());
    return input_ptr; // 返回智能指针
}

void LLM::chat_inputs(struct ChatConfig config,
                      std::vector<tc::InferInput *> &inputs) {
    auto ptr_request_type =
        createInput("request_type", {1}, "INT32", config.request_type_data);
    inputs.push_back(ptr_request_type.get());
    auto ptr_request_id =
        createInput("request_id", {1}, "BYTES", std::vector{config.request_id});
    inputs.push_back(ptr_request_id.get());
    auto ptr_text_input =
        createInput("text_input", {1}, "BYTES", std::vector{config.chatPrompt});
    inputs.push_back(ptr_text_input.get());
    auto ptr_temperature_input = createInput("temperature_input", {1}, "FP32",
                                             config.temperature_input_data);
    inputs.push_back(ptr_temperature_input.get());
    auto ptr_top_k_input =
        createInput("top_k_input", {1}, "INT32", config.top_k_input_data);
    inputs.push_back(ptr_top_k_input.get());
    auto ptr_top_p_input =
        createInput("top_p_input", {1}, "FP32", config.top_p_input_data);
    inputs.push_back(ptr_top_p_input.get());
    auto ptr_n_keep_input =
        createInput("n_keep_input", {1}, "INT32", config.n_keep_input_data);
    inputs.push_back(ptr_n_keep_input.get());
    auto ptr_n_predict_input = createInput("n_predict_input", {1}, "INT32",
                                           config.n_predict_input_data);
    inputs.push_back(ptr_n_predict_input.get());
    auto ptr_cache_prompt_input = createInput("cache_prompt_input", {1}, "BOOL",
                                              config.cache_prompt_input_data);
    inputs.push_back(ptr_cache_prompt_input.get());
    auto ptr_stop_input = createInput(
        "stop_input", {static_cast<int64_t>(config.stop_input_data.size())},
        "BYTES", config.stop_input_data);
    inputs.push_back(ptr_stop_input.get());
    auto ptr_lora_scale_input =
        createInput("lora_scale_input",
                    {static_cast<int64_t>(config.lora_scale_input_data.size())},
                    "FP32", config.lora_scale_input_data);
    inputs.push_back(ptr_lora_scale_input.get());
    auto ptr_stream_input =
        createInput("stream_input", {1}, "BOOL", config.stream_input_data);
    inputs.push_back(ptr_stream_input.get());

    // Create text output
    tc::InferRequestedOutput *text_output;
    FAIL_IF_ERR(tc::InferRequestedOutput::Create(&text_output, "text_output"),
                "unable to get text_output");
    std::shared_ptr<tc::InferRequestedOutput> text_output_ptr(text_output);

    // Async inference call
    tc::InferOptions options(modelName_);
    options.model_version_ = "";

    tc::Error err = inferClient_->AsyncStreamInfer(options, inputs,
                                                   {text_output_ptr.get()});
    LOGI("Stream Inference: Start Inference--- {}", config.request_id);

    if (!err.IsOk()) {
        LOGE("{} AsyncStreamInfer error: {}", __LINE__, err.Message());
    }
}
#endif
std::string LLM::chat(const std::string &prompt, const int &task_id,
                      struct ChatConfig &config) {
#ifdef TRITON_CLIENT
    is_chatting_.store(true, std::memory_order_release);
    bool verbose = true;

    result_list_.clear();

    // int lora_count = getLoraCount();
    int lora_count =
        2; // 为规避triton接收无lora的推理请求会报错的问题做出的临时修改
    std::vector<float> lora_scale_input_data(lora_count, 0.0f);
    setLoraScale(task_id, lora_scale_input_data);

    config.request_id = request_id_;
    config.chatPrompt = prompt;
    config.n_predict_input_data = n_predict_input_;
    config.lora_scale_input_data = lora_scale_input_data;
    config.stream_input_data = true;
    std::vector<tc::InferInput *> inputs;

    std::string ans;
    bool finalResponse = false;

    chat_inputs(config, inputs);
    bool hasReceivedToken = false;
    int recv_count = 0;
    while (!finalResponse) {
        // Wait until all callbacks are invoked
        {
            std::unique_lock<std::mutex> lk(mutex_);
            cv_.wait(lk, [&]() {
                if (result_list_.size() > 0) {
                    return true;
                } else {
                    return false;
                }
            });
            if (!hasReceivedToken) {
                LOGI("Stream Inference: get first token --- {}",
                     config.request_id);
                hasReceivedToken = true;
            }
            std::shared_ptr<tc::InferResult> result = result_list_.front();
            result->IsFinalResponse(&finalResponse);
            std::vector<std::string> text_output_data;
            result->StringData("text_output", &text_output_data);

            if (text_output_data.size() != 1) {
                LOGE("error: received incorrect number of strings for "
                     "text_output:  {}",
                     text_output_data.size());
                break;
            }
            ans += text_output_data[0];
            printf("prompt:%s\n<----->nlp engine get from triton llm:%s\n",
                   prompt.c_str(), ans.c_str());
            if (config.stream_output) {
                completionCallbacks_(text_output_data[0], false);
            }
            result_list_.erase(result_list_.begin());
        }
    }
    LOGI("Stream Inference: get last token --- {}", config.request_id);
    std::string end_flag = "";

    if (!config.stream_output) {
        completionCallbacks_(ans, false);
    }
    completionCallbacks_(end_flag, true);
    is_chatting_.store(false, std::memory_order_release);

    return ans;
#else
    completionCallbacks_("", true);
    return "";
#endif
}
