dpplogger/main.cpp

#include "Random.hpp"
#include "sqlite_modern_cpp/sqlite_modern_cpp.h"

#include <iostream>
#include <unordered_map>
#include <optional>
#include <chrono>
#include <memory>
#include <stdexcept>
#include <boost/asio/io_service.hpp>
#include <dcboost/discord.hpp>
#include <dcboost/snowflake.hpp>

using namespace ChatFuse;
using namespace QLog;



class Cache {
    struct Miss : public std::runtime_error {
        Miss() : std::runtime_error("Object cache miss") {}
    };

    Logger<Cache> logger;
    std::unordered_map<Discord::Snowflake, Json::Value> cache;

public:
    Cache() {
        logger.inst_ptr = this;
    }

    const Json::Value& store(const Json::Value& object) {
        const auto& id_str = object["id"].asString();
        logger.log(Loglevel::verbose, "Stored "+id_str+" in cache");
        return cache.insert_or_assign(std::move(id_str), object).first->second;
    }
    const Json::Value& fetch(Discord::Snowflake id) const {
        const auto& entry = cache.find(id);
        if (entry == cache.end()) {
            logger.log(Loglevel::warn, "Missing "+id.str());
            throw Miss();
        }
        return entry->second;
    }
    bool has(Discord::Snowflake id) const {
        return cache.find(id) != cache.end();
    }
};


struct MyClient final : public Discord::Client {
    sqlite::database db;
    Cache cache;

    Json::Value *me;

public:
    MyClient(boost::asio::io_service& io, const ChatFuse::Discord::Settings& settings = {})
        : ChatFuse::Discord::Client(io, settings),
          db("log.sqlite3")
    {
        // Create tables
        {
            db << "CREATE TABLE IF NOT EXISTS messages ("
                  "    id TEXT PRIMARY KEY NOT NULL,"
                  "    type INTEGER NOT NULL,"
                  "    channel_id TEXT NOT NULL,"
                  "    author_id TEXT NOT NULL,"
                  "    replied_to_id TEXT,"
                  "    is_deleted INTEGER DEFAULT 0 NOT NULL,"
                  "    is_edited INTEGER DEFAULT 0 NOT NULL,"
                  "    creation_timestamp TEXT NOT NULL,"
                  "    update_timestamp TEXT DEFAULT '0' NOT NULL,"
                  "    UNIQUE(id)"
                  ");";
            db << "CREATE TABLE IF NOT EXISTS message_contents ("
                  "    message_id TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    content TEXT NOT NULL"
                  ");";
            db << "CREATE TABLE IF NOT EXISTS users ("
                  "    id TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    full_username TEXT NOT NULL," // Username#Discriminator
                  "    avatar TEXT,"
                  "    bio TEXT,"
                  "    has_nitro INTEGER NOT NULL,"
                  "    is_bot INTEGER NOT NULL"
                  ");";
            db << "CREATE TABLE IF NOT EXISTS user_statuses ("
                  "    user_id TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    status TEXT," // JSON array or null if not updated: [int:online/dnd/afk/offline enum, ?:emoji (null/empty if none), string:text (empty if none)]
                  "    presence TEXT" // Discord JSON presence object
                  ");";
            db << "CREATE TABLE IF NOT EXISTS members ("
                  "    user_id TEXT NOT NULL,"
                  "    guild_id TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    nickname TEXT,"
                  "    in_guild INTEGER NOT NULL,"
                  "    is_untracked INTEGER NOT NULL"
                  ");";
            db << "CREATE TABLE IF NOT EXISTS member_voice_connections ("
                  "    channel_id TEXT NOT NULL,"
                  "    user_id TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    voice_channel_id TEXT,"
                  "    self_muted INTEGER NOT NULL,"
                  "    self_deafened INTEGER NOT NULL,"
                  "    server_muted INTEGER NOT NULL,"
                  "    server_deafened INTEGER NOT NULL"
                  ");";
            db << "CREATE TABLE IF NOT EXISTS channels ("
                  "    id TEXT NOT NULL,"
                  "    guild_id TEXT NOT NULL,"
                  "    category_channel_id TEXT NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    type INTEGER NOT NULL,"
                  "    name TEXT NOT NULL,"
                  "    topic TEXT,"
                  "    has_access INTEGER NOT NULL"
                  ");";
            db << "CREATE TABLE IF NOT EXISTS guilds ("
                  "    id TEXT NOT NULL,"
                  "    timestamp TEXT NOT NULL,"
                  "    is_initial INTEGER DEFAULT 1 NOT NULL,"
                  "    name TEXT NOT NULL,"
                  "    owner_user_id TEXT NOT NULL,"
                  "    in_guild INTEGER DEFAULT 1 NOT NULL"
                  ");";
        }
    }

    boost::asio::awaitable<void> fetchAllGuilds(const Json::Value& unavailable_guild_array) {
        RandomGenerator rng;
        rng.seed();
        for (const auto& incomplete_data : unavailable_guild_array) {
            Discord::Snowflake guild_id = incomplete_data["id"];

            // Make sure guild isn't cached yet
            if (cache.has(guild_id)) continue;

            // Make sure data is actually incomplete
            if (incomplete_data["name"].isString()) {
                // We'll just use it as-is
                cache.store(incomplete_data);
                insertGuildUpdate(incomplete_data, true);
                continue;
            }

            // Fetch guild from API and store it in cache
            auto data = co_await api.call(boost::beast::http::verb::get, "/guilds/"+guild_id.str());
            cache.store(data);

            // Insert guild into database
            insertGuildUpdate(data, true);

            // Delay
            co_await asyncSleep(rng.getUInt(6000, 15000));
        }
    }

    void insertGuildUpdate(const Json::Value& data, bool is_initial) {
        db << "INSERT OR IGNORE INTO guilds (id, timestamp, is_initial, name, owner_user_id)"
              "                      VALUES (?,  ?,         ?,          ?,    ?            );"
           << data["id"].asString() << std::to_string(time(nullptr)) << is_initial << data["name"].asString() << data["owner_id"].asString();
        cache.store(data);
    }
    void insertGuildLeave(const Json::Value& data) {
        auto cached_data = cache.fetch(data["id"]);
        db << "INSERT OR IGNORE INTO guilds (id, timestamp, is_initial, name, owner_user_id, in_guild)"
              "                      VALUES (?,  ?,         0,          ?,    ?,             0       );"
           << data["id"].asString() << std::to_string(time(nullptr)) << cached_data["name"].asString() << cached_data["owner_id"].asString();
    }

    virtual boost::asio::awaitable<void> intentHandler(const std::string& intent, const Json::Value& data) override {
        if (intent == "READY") [[unlikely]] {
            const auto& user = cache.store(data["user"]);
            logger.log(Loglevel::info, "Connected to Discord as: "+user["username"].asString()+'#'+user["discriminator"].asString());
            settings.is_bot = user["bot"].asBool();
            co_await fetchAllGuilds(data["guilds"]);
        } else if (intent == "GUILD_CREATE") [[unlikely]] {
            insertGuildUpdate(data, true);
        } else if (intent == "GUILD_UPDATE") [[unlikely]] {
            insertGuildUpdate(data, false);
        } else if (intent == "GUILD_DELETE") [[unlikely]] {
            insertGuildLeave(data);
        }
        co_return;
    }
};


int main(int argc, char **argv) {
    // Check args
    if (argc == 1) {
        std::cout << "Usage: " << argv[0] << " <token>" << std::endl;
        return EXIT_FAILURE;
    }

    // Create io service
    boost::asio::io_service io;

    // Set up client
    auto client = std::make_shared<MyClient>(io, ChatFuse::Discord::Settings{
                                                 .bot_token = argv[1],
                                                 .intents = ChatFuse::Discord::intents::all
                                             });
    client->detach();

    // Erase bot token from argv (so it's no longer visible in /proc/{pid}/cmdline)
    memset(argv[1], 0, strlen(argv[1]));

    // Run!!!
    io.run();

    return EXIT_SUCCESS;



    /*// Initialize bot
    dpp::cluster cluster("ODAyODYzNTU0MjY4NjI2OTY1.GsYaeE.Mzu8hrICXwcJMfDKnYssRMkRtDfX-r7zrQOCRI", dpp::i_all_intents);

    // Set up callbacks
    // Useful: https://discord.com/developers/docs/topics/gateway-events
    cluster.on_ready([&](const dpp::ready_t& event) {
        std::cout << "Connected to Discord as " << cluster.me.format_username() << std::endl;
    });

    cluster.on_guild_create([&](const dpp::guild_create_t& event) {
        std::cout << "Logging guild " << event.created->name << std::endl;
        const auto guild = event.created;
        db << "INSERT OR IGNORE INTO levels (id, timestamp, name, owner_user_id)"
              "                      VALUES (?,  ?,         ?,    ?            );"
           << std::to_string(guild->id) << std::to_string(time(nullptr)) << guild->name << std::to_string(guild->owner_id);
        cache.store(guild->id, *guild);
    });
    cluster.on_guild_update([&](const dpp::guild_update_t& event) {
        const auto guild = event.updated;
        db << "INSERT OR IGNORE INTO levels (id, timestamp, is_initial, name, owner_user_id)"
              "                      VALUES (?,  ?,         0,          ?,    ?            );"
           << std::to_string(guild->id) << std::to_string(time(nullptr)) << guild->name << std::to_string(guild->owner_id);
        cache.store(guild->id, *guild);
    });
    cluster.on_guild_delete([&](const dpp::guild_delete_t& event) {
        const auto& guild = std::get<dpp::guild>(cache.fetch(event.deleted->id));
        db << "INSERT OR IGNORE INTO levels (id, timestamp, is_initial, name, owner_user_id)"
              "                      VALUES (?,  ?,         0,          ?,    ?            );"
           << std::to_string(guild.id) << std::to_string(time(nullptr)) << guild.name << std::to_string(guild.owner_id);
    });

    auto on_message_content = [&](const dpp::message& msg, bool is_initial) {
        db << "INSERT OR IGNORE INTO message_contents (message_id, is_initial, timestamp, content)"
              "                                VALUES (?,          ?,          ?,         ?      );"
           << std::to_string(msg.id) << is_initial << std::to_string(msg.edited?msg.edited:msg.sent) << msg.content;
    };
    cluster.on_message_create([&](const dpp::message_create_t& event) {
        const auto& msg = event.msg;
        db << "INSERT OR IGNORE INTO messages (id, type, channel_id, author_id, replied_to_id, creation_timestamp)"
              "                        VALUES (?,  ?,    ?,          ?,         ?,             ?                 );"
           << std::to_string(msg.id) << std::to_string(msg.type) << std::to_string(msg.channel_id) << std::to_string(msg.author.id) << (msg.message_reference.message_id?std::optional<std::string>(std::to_string(msg.message_reference.message_id)):nullptr) << std::to_string(msg.sent);
        if (!msg.content.empty()) on_message_content(msg, true);
        cache.store(msg.author.id, msg.author);
    });
    cluster.on_message_update([&](const dpp::message_update_t& event) {
        const auto& msg = event.msg;
        if (!msg.content.empty()) {
            on_message_content(msg, false);
            db << "UPDATE messages "
                  "SET is_edited = 1 "
                  "WHERE id = ?;"
               << std::to_string(msg.id);;
        }
    });
    cluster.on_message_delete([&](const dpp::message_delete_t& event) {
        db << "UPDATE messages "
              "SET is_deleted = 1 "
              "WHERE id = ?;"
           << std::to_string(event.deleted->id);;
    });

    // Run
    cluster.start(false);*/
}