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

#ifndef SERVICES_CHATBOT_H
#define SERVICES_CHATBOT_H

#include <kylin-ai/aiengine.h>
#include <modelconfig/modelconfigenums.h>

#include <memory>

#include "chatmessages.h"
#include "concurrentqueue/blockingconcurrentqueue.h"
#include "config/prompttemplate.h"
#include "engineselector/engineselector.h"
#include "taskmanager/taskrunner.h"

enum class TaskId {
    NormalChat = -1,
    TagExtraction,
    IntentRecognition,
};

class Chatbot : public std::enable_shared_from_this<Chatbot> {
private:
    using INlpEngine = ai_engine::gen_ai::nlp::AbstractTextGenerationEngine;

public:
    struct ChatResult {
        std::string reason_;
        std::string reasoningContent_;
        std::string modelName_;
        std::string finishReason_;
        ai_engine::EngineError error_;
        bool isEnd_{true};
    };

    using ChatResultCallback = std::function<void(const ChatResult &)>;
    using ChatFinishedCallback = std::function<void()>;

public:
    Chatbot(EngineSelector &engineSelector, bool isDefaultConfig);

    ~Chatbot();

    void chat(const std::string &question, ChatResultCallback chatResultCallback,
              ChatFinishedCallback finishedCallback);

    void chatWithHistoryMessage(const std::string &messages, ChatResultCallback chatResultCallback,
                                ChatFinishedCallback finishedCallback);

    void stopChat();

    void clearChatContext();

    void setTopK(int topK) { chatMessages_.setTopK(topK); }

    void addSystemMessage(const std::string &message);

    void addSystemMessage(const std::string &message, int taskId);

    void addSystemMessageByPromptId(int promptId, DeployType deployType);

    void addUserMessage(const std::string &message);

    void addAssistantMessage(const std::string &message);

    void setDeployType(DeployType deployType) { deployType_ = deployType; }

    DeployType getDeployType() const { return deployType_; }

    // 返回具体的错误码
    [[nodiscard]] int init(const std::shared_ptr<INlpEngine> &nlpEngine);

    void destroy();

private:
    struct ChatRequest {
        std::string question;
        bool withHistoryMessage{false};
        ChatResultCallback resultCallback{nullptr};
        ChatFinishedCallback finishedCallback{nullptr};
    };

private:
    std::shared_ptr<Chatbot> getSharedObject() { return shared_from_this(); }
    void handleChatResult(const ai_engine::gen_ai::nlp::TextGenerationResult &result);
    void callChatResultCallback(const ChatResult &result) const;
    void callChatFinishedCallback() const;

    void chatAsync();

    void chatTask();

    bool handleNetworkDisconnection(const std::string &message);

    DeployType currentEngineDeployType() const;

private:
    EngineSelector &engineSelector_;
    std::shared_ptr<INlpEngine> nlpEngine_{nullptr};
    std::shared_ptr<config::PromptTemplate> promptTemplate_{nullptr};
    ChatMessages chatMessages_;
    std::string replyCache_;
    DeployType deployType_{DeployType::Default};
    ChatRequest currentChatRequest_;
    moodycamel::BlockingConcurrentQueue<ChatRequest> chatRequestQueue_;
    std::atomic<task::TaskId> taskId_{task::INVALID_TASK_ID};
    std::atomic_bool isInited_{false};
    std::atomic_bool isDefaultConfig_{true};
};

#endif  // SERVICES_CHATBOT_H
