shard_mgr.hpp

//
// shard_mgr.hpp
// *************
//
// Copyright (c) 2019 Sharon W (sharon at aegis dot gg)
//
// Distributed under the MIT License. (See accompanying file LICENSE)
// 

#pragma once

#include "aegis/config.hpp"
#include "aegis/utility.hpp"
#include "aegis/snowflake.hpp"
#include "aegis/error.hpp"
#include "aegis/shards/shard.hpp"
#include <websocketpp/config/asio_client.hpp>
#include <websocketpp/common/connection_hdl.hpp>
#include <websocketpp/client.hpp>
#include "spdlog/spdlog.h"

#include <vector>
#include <iostream>
#include <string>

#include <asio/bind_executor.hpp>

namespace aegis
{

namespace shards
{

using json = nlohmann::json;
using namespace std::literals;


class shard_mgr
{
public:
    using websocket = websocketpp::client<websocketpp::config::asio_tls_client>;

    using connection_ptr = websocketpp::client<websocketpp::config::asio_tls_client>::connection_type::ptr;

    using message_ptr = websocketpp::config::asio_client::message_type::ptr;


    AEGIS_DECL shard_mgr(std::string token, asio::io_context & _io, std::shared_ptr<spdlog::logger> log);

    AEGIS_DECL ~shard_mgr();

    shard_mgr(const shard_mgr &) = delete;
    shard_mgr(shard_mgr &&) = delete;
    shard_mgr & operator=(const shard_mgr &) = delete;


    AEGIS_DECL void setup_callbacks(shard * _shard) noexcept;


    AEGIS_DECL void debug_trace(shard * _shard, bool extended = false) noexcept;


    asio::io_context & get_io_context()
    {
        return _io_context;
    }

    AEGIS_DECL void shutdown();


    websocket & get_websocket() noexcept { return websocket_o; }


    bot_status get_state() const noexcept { return _status; }


    void set_state(bot_status s) noexcept { _status = s; }


    AEGIS_DECL std::string uptime() const noexcept;


    AEGIS_DECL void send_all_shards(const std::string & msg);


    AEGIS_DECL void send_all_shards(const json & msg);

    AEGIS_DECL void start();

    using t_on_message = std::function<void(websocketpp::connection_hdl hdl, std::string msg, shard * _shard)>;
    using t_on_connect = std::function<void(websocketpp::connection_hdl hdl, shard * _shard)>;
    using t_on_close = std::function<void(websocketpp::connection_hdl hdl, shard * _shard)>;


    void set_on_message(t_on_message cb) noexcept
    {
        i_on_message = cb;
    }


    void set_on_connect(t_on_connect cb) noexcept
    {
        i_on_connect = cb;
    }


    void set_on_close(t_on_close cb) noexcept
    {
        i_on_close = cb;
    }


    void set_gateway_url(const std::string & url) noexcept
    {
        gateway_url = url;
    }


    std::string get_gateway_url() const noexcept
    {
        return gateway_url;
    }


    AEGIS_DECL void reset_shard(shard * _shard, shard_status _status = shard_status::closing) noexcept;


    AEGIS_DECL void queue_reconnect(shard * _shard) noexcept;


    void queue_reconnect(shard & _shard)
    {
        queue_reconnect(&_shard);
    }


    AEGIS_DECL void connect(shard * _shard) noexcept;


    AEGIS_DECL shard & get_shard(uint16_t shard_id);


    AEGIS_DECL const std::vector<std::unique_ptr<shard>> & get_shards() const noexcept
    {
        return _shards;
    }


    AEGIS_DECL void close(shard * _shard, int32_t code = 1001, const std::string & reason = "", shard_status connection_state = shard_status::closing) noexcept;


    void close(shard & _shard, int32_t code = 1001, const std::string & reason = "", shard_status connection_state = shard_status::closing) noexcept
    {
        close(&_shard, code, reason, connection_state);
    }


    uint32_t shard_count() const noexcept
    {
        return _shards.size();
    }

    asio::io_context & _io_context;

    std::string ws_gateway;

    std::unordered_map<std::string, uint64_t> message_count;

    uint32_t force_shard_count;
    uint32_t shard_max_count;
    std::shared_ptr<spdlog::logger> log;

private:
    friend aegis::core;

    std::chrono::time_point<std::chrono::steady_clock> _last_ready;
    std::chrono::time_point<std::chrono::steady_clock> _last_identify;
    std::chrono::time_point<std::chrono::steady_clock> _connect_time;
    shard * _connecting_shard;

    std::vector<std::unique_ptr<shard>> _shards;
  
    t_on_message i_on_message;
    t_on_connect i_on_connect;
    t_on_close i_on_close;

    std::function<void(aegis::shards::shard*)> i_shard_disconnect;
    std::function<void(aegis::shards::shard*)> i_shard_connect;

    AEGIS_DECL void _on_message(websocketpp::connection_hdl hdl, message_ptr msg, shard * _shard);
    AEGIS_DECL void _on_connect(websocketpp::connection_hdl hdl, shard * _shard);
    AEGIS_DECL void _on_close(websocketpp::connection_hdl hdl, shard * _shard);
    AEGIS_DECL void ws_status(const asio::error_code & ec);

    std::chrono::steady_clock::time_point starttime;

    // Gateway URL for the Discord Websocket
    std::string gateway_url;

    // Websocket++ object
    websocket websocket_o;

    // Bot's token
    std::string _token;

    bot_status _status;

    std::shared_ptr<asio::steady_timer> ws_timer;
    std::shared_ptr<asio::steady_timer> ws_connect_timer;
    std::deque<shard*> _shards_to_connect;
};

}

}

#if defined(AEGIS_HEADER_ONLY)
#include "aegis/shards/impl/shard_mgr.cpp"
#endif