#include "NDSUI.hpp"
#include "Runtime.hpp"
#include "AsyncManager.hpp"
#include "Client.hpp"
#include "basic-coro/AwaitableTask.hpp"

#include <iostream>
#include <string>
#include <string_view>



void on_progress(float progress) {
    std::cout << unsigned(progress) << '\r' << std::flush;
}

basiccoro::AwaitableTask<void> async_main(Runtime& rt, AsyncManager &aMan, NDSUI& ui) {
    // Ask for server address
    ui.addLogMessage(ui.createLogMessage("System", "Please type in the address of the server."));
    const std::string addr = co_await ui.readLine();

    // Create client
    Client client(addr, 99181, aMan);

    // Greet the user
    ui.addLogMessage(ui.createLogMessage("Bot", "Hey! How can I help you today?"));

    // Connection loop
    for (;; rt.cooperate()) {
        // Read prompt
        const auto prompt = co_await ui.readLine();

        // Display prompt
        auto& msg = ui.addLogMessage(ui.createLogMessage("Bot", "Initializing..."));

        // Hide keyboard
        keyboardHide();

        // Clear console
        consoleClear();

        // Run inference
        std::string result;
        co_await client.ask(prompt, [&msg, &ui] (unsigned progress) -> basiccoro::AwaitableTask<void> {
            msg = ui.createLogMessage("Bot", (std::to_string(progress)+"%"));
            // Check if ready to generate
            if (progress == 100) {
                // Show help
                std::cout << "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n"
                             "Stop now:         [START] Abort \n"
                             "Stop soon:       [SELECT] Cancel" << std::flush;
            }
            co_return;
        }, [&result, &msg, &ui] (std::string_view token) -> basiccoro::AwaitableTask<void> {
            // Update message with current generation result
            result.append(token);
            msg = ui.createLogMessage("Bot", result+"...");
            co_return;
        });
        msg = ui.createLogMessage("Bot", result);

        // Clear console
        consoleClear();

        // Show keyboard
        keyboardShow();
    }
}

int main() {
    Runtime rt;
    NDSUI ui;
    AsyncManager aMan(rt);

    // Print header
    std::cout << "llama.nds running on " PLATFORM ".\n"
                 "\n";

    // Start async main()
    async_main(rt, aMan, ui);

    // Start async manager
    while (aMan.shouldRun() || ui.shouldRun()) {
        ui.run();
        aMan.run();
    }

    return 0;
}