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

#include <jsoncpp/json/value.h>

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

std::string ChatMessages::generateJsonString() {
    Json::Value messages(Json::arrayValue);

    Json::Value systemMessage(Json::objectValue);
    if (!systemMessage_.empty()) {
        systemMessage["role"] = "system";
        systemMessage["content"] = systemMessage_;
        messages.append(systemMessage);
    }

    for (const auto& item : messages_) {
        Json::Value message(Json::objectValue);
        if (item.type_ == ChatMessageType::USER) {
            message["role"] = "user";
            message["content"] = item.message_;
        } else if (item.type_ == ChatMessageType::ASSISTANT) {
            message["role"] = "assistant";
            message["content"] = item.message_;
        } else {
            // do nothing
        }
        messages.append(message);
    }

    Json::Value requestBody(Json::objectValue);
    requestBody["messages"] = messages;
    if (topK_ != 0) {
        requestBody["top_k"] = topK_;
    }

    if (taskId_ != -1) {
        requestBody["task_id"] = taskId_;
    }
    std::string ret = util::jsonToString(requestBody);
    LOGD("Request body: {}", ret);
    return ret;
}

void ChatMessages::setSystemMessage(const std::string& systemMessage) { systemMessage_ = systemMessage; }

void ChatMessages::setSystemMessage(const std::string& systemMessage, int taskId) {
    taskId_ = taskId;
    systemMessage_ = systemMessage;
}

void ChatMessages::addUserQuestion(const std::string& question) {
    if (question.empty()) {
        return;
    }
    messages_.push_back({ChatMessageType::USER, question});
}

void ChatMessages::addAssistantReply(const std::string& reply) {
    if (reply.empty()) {
        return;
    }
    messages_.push_back({ChatMessageType::ASSISTANT, reply});
}