/*
 * 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/>.
 */

#pragma once

#include <condition_variable>
#include <string>
#ifdef TRITON_CLIENT
#include <triton/client/grpc_client.h>
#include <triton/client/http_client.h>
#endif
#include "../utils/logger.h"

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

using CompletionResultCallback =
    std::function<void(const std::string &, bool over)>;
#ifdef TRITON_CLIENT
using ResultList = std::vector<std::shared_ptr<tc::InferResult>>;
#endif

struct ChatConfig {
    int request_type_data = 0;
    std::string request_id = "message_1";
    std::string chatPrompt = "message";
    float temperature_input = 0.0;
    int32_t temperature_input_data = 0.0;
    int32_t top_k_input_data = 40;
    float top_p_input_data = 0.95;
    int32_t n_keep_input_data = 0;
    int32_t n_predict_input_data = 512;
    bool cache_prompt_input_data = true;
    std::vector<std::string> stop_input_data{"<|im_end|>"};
    std::vector<float> lora_scale_input_data = {0, 0};
    bool stream_input_data = true;
    bool stream_output = true;
};

enum class ClientPreparationStatus {
    NotAttempted, // 未尝试
    Success,      // 成功
    Failed        // 失败
};

static bool modelBankIsReady_ = false;
class LLM {
public:
    LLM();
    ~LLM();
    bool prepareModelInfo();
    int setLoraScale(int task_id, std::vector<float> &lora_scale_input_data);
    std::string chat(const std::string &prompt, const int &task_id,
                     struct ChatConfig &config);
    void setCompletionResultCallback(CompletionResultCallback callback);
    std::string GetConfigValue(const char *key, const char *engine);
    bool readySetLlmModelName(const std::string &modelName);
    bool doSetLlmModelName();
    void stopChat();
    int getLoraCount();
#ifdef TRITON_CLIENT
    void llm_callback(tc::InferResult *result);
#endif
    bool llm_load_model();
    void modelInfoUpdate();
    std::mutex mutex_;
    std::condition_variable cv_;
    std::string targetModelName_ = "";
    std::string modelName_;
    std::string modelDisplayName_ = "";
    std::string modelInfo_;
#ifdef TRITON_CLIENT
    ResultList result_list_;
#endif
    int n_predict_input_;
    ClientPreparationStatus clientPreparationStatus_ =
        ClientPreparationStatus::NotAttempted;

private:
#ifdef TRITON_CLIENT
    void chat_inputs(struct ChatConfig config,
                     std::vector<tc::InferInput *> &inputs);
#endif
    std::atomic<bool> is_chatting_{false};
    std::string request_id_ = "";
    std::string modelVersion_ = "";
    std::string default_model_info_ = "";
    std::string grpcServerUrl_ = "localhost:8001";
    int modelLoadTimeoutUs_ = 60 * 1000 * 1000;
    CompletionResultCallback completionCallbacks_;
#ifdef TRITON_CLIENT

    std::unique_ptr<tc::InferenceServerGrpcClient> loadClient_ = nullptr;
    std::unique_ptr<tc::InferenceServerGrpcClient> inferClient_ = nullptr;
#endif
};