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

// std
#include <stdexcept>
#include <string>
#include <vector>

// api
#include "kylin-ai/genai/text/chat.h"
#include "kylin-ai/genai/text/chatmessage.h"
#include "kylin-ai/genai/text/config.h"
#include "kylin-ai/genai/text/error.h"
#include "kylin-ai/genai/text/result.h"

// glib
#include <gio/gio.h>

namespace {
std::vector<std::string> stringSplit(const std::string &str,
                                     const std::string &delim) {
    std::vector<std::string> result;
    size_t begin = 0;
    for (size_t end; end = str.find(delim, begin), end != std::string::npos;
         begin = end + delim.size()) {
        result.emplace_back(str.substr(begin, end - begin));
    }
    result.emplace_back(str.substr(begin));
    return result;
}
}  // namespace

class Chatbot {
public:
    Chatbot() {
        stdinChannel_ = g_io_channel_unix_new(STDIN_FILENO);
        g_io_channel_set_encoding(stdinChannel_, nullptr, nullptr);
        g_io_channel_set_buffered(stdinChannel_, TRUE);
        g_io_channel_set_flags(stdinChannel_, G_IO_FLAG_NONBLOCK, nullptr);
        g_io_add_watch(stdinChannel_,
                       GIOCondition(G_IO_IN | G_IO_HUP | G_IO_ERR),
                       onStdinEnter, this);

        mainLoop_ = g_main_loop_new(nullptr, false);
        session_ = genai_text_create_session();

        config_ = chat_model_config_create();

#define BAIDU 0
#if BAIDU
        const char *MODEL_NAME = "百度-ERNIE-Bot-4";

        chat_model_config_set_name(config_, MODEL_NAME);
        chat_model_config_set_top_k(config_, 4);
        chat_model_config_set_deploy_type(config_,
                                          ModelDeployType::PublicCloud);

        genai_text_set_model_config(session_, config_);
#else
        // default
        genai_text_set_model_config(session_, nullptr);
#endif

        int errorCode = genai_text_init_session(session_);
        if (errorCode != AISDK_NO_ERROR) {
            std::string errorMessage =
                "Func genai_text_init_session init error, error code is: " +
                std::to_string(errorCode);
            throw std::runtime_error(errorMessage.c_str());
        }
        genai_text_result_set_callback(session_, callback, nullptr);
    }

    ~Chatbot() {
        g_main_loop_unref(mainLoop_);
        g_io_channel_unref(stdinChannel_);
    }

    void run() const { g_main_loop_run(mainLoop_); }

private:
    static gboolean onStdinEnter(GIOChannel *channel, GIOCondition condition,
                                 gpointer userData) {
        Chatbot *tool = (Chatbot *)userData;
        gchar *line{nullptr};
        if (condition & G_IO_IN) {
            GIOStatus status = g_io_channel_read_line(channel, &line, nullptr,
                                                      nullptr, nullptr);
            if (status == G_IO_STATUS_NORMAL) {
                std::string input(line);
                if (!input.empty() && input[input.size() - 1] == '\n') {
                    input.erase(input.size() - 1);
                }
                tool->control(input);
                g_free(line);
            }
        }
        return TRUE;
    }

    void control(const std::string &input) const {
        if (input.empty()) {
            std::fprintf(stdout, "Input message is empty.\n");
        } else if (input == "quit") {
            g_main_loop_quit(mainLoop_);
        } else if (input.substr(0, 5) == "top_k") {
            std ::vector<std::string> args = stringSplit(input, " ");
            int topK = std::stoi(args[1]);
            chat_model_config_set_top_k(config_, topK);
            genai_text_set_model_config(session_, config_);
        } else {
            std::fprintf(stdout, "----------------------------------------\n");
            std::fprintf(stdout, "question:\n");
            std::fprintf(stdout, "%s\n", input.c_str());
            std::fprintf(stdout, "answer:\n");
            genai_text_chat_async(session_, input.c_str());
        }
    }

    static void callback(ChatResult *result, void *user_data) {
        std::fprintf(stdout, "%s", chat_result_get_assistant_message(result));
        if (chat_result_get_is_end(result)) {
            std::fprintf(stdout, "\n");
            std::fprintf(stdout, "========================================\n");
        }
    }

private:
    GIOChannel *stdinChannel_{nullptr};
    GMainLoop *mainLoop_{nullptr};

    ChatModelConfig *config_{nullptr};
    GenAiTextSession *session_{nullptr};
};

int main() {
    try {
        Chatbot chatbot{};
        chatbot.run();
    } catch (std::exception &exception) {
        std::fprintf(stderr, "Exception: %s\n", exception.what());
        return -1;
    }
    return 0;
}