#include <set>
#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/server.hpp>

typedef websocketpp::server<websocketpp::config::asio> server;

using websocketpp::connection_hdl;
using websocketpp::lib::placeholders::_1;
using websocketpp::lib::placeholders::_2;
using websocketpp::lib::bind;

struct connection_data {
    int sessionid;
    bool hello;
};

class broadcast_server {
public:
    uint8_t log_level = 0;

    broadcast_server() : m_next_sessionid(1) {
        m_server.init_asio();

        m_server.set_open_handler(bind(&broadcast_server::on_open, this, ::_1));
        m_server.set_close_handler(bind(&broadcast_server::on_close, this, ::_1));
        m_server.set_message_handler(bind(&broadcast_server::on_message, this, ::_1, ::_2));
    }

    void on_open(connection_hdl hdl) {
        connection_data data;

        data.sessionid = m_next_sessionid++;
        data.hello = false;

        m_connections[hdl] = data;
    }

    void on_close(connection_hdl hdl) {
        m_connections.erase(hdl);
    }

    void on_message(connection_hdl hdl, server::message_ptr msg) {
        connection_data& data = get_data_from_hdl(hdl);

        if (msg->get_opcode() == websocketpp::frame::opcode::binary) {
            data.hello = true;
            m_server.send(hdl, msg);
        } else {
            this->close(hdl, websocketpp::close::status::normal);
        }

        /*
        for (auto it : m_connections) {
            // m_server.send(it, msg);
        }
        */
    }

    void close(connection_hdl hdl, websocketpp::close::status::value code) {
        websocketpp::lib::error_code ec;

        m_server.close(hdl, code, "", ec);

        if (ec && log_level > 0) {
            std::cout << "> Error initiating close: " << ec.message() << std::endl;
        }
    }

    void log(uint8_t level) {
        log_level = level;
        m_server.clear_access_channels(websocketpp::log::alevel::all);

        switch (level) {
            case 1:
                m_server.set_access_channels(websocketpp::log::alevel::connect | websocketpp::log::alevel::disconnect);
                break;
            case 2:
                m_server.set_access_channels(websocketpp::log::alevel::connect | websocketpp::log::alevel::disconnect);
                m_server.set_access_channels(websocketpp::log::alevel::frame_header | websocketpp::log::alevel::frame_payload);
                break;
        }
    }

    void run(uint16_t port) {
        if (log_level == 0) m_server.clear_access_channels(websocketpp::log::alevel::all);
        m_server.listen(port);
        m_server.start_accept();
        m_server.run();
    }

    connection_data& get_data_from_hdl(connection_hdl hdl) {
        auto it = m_connections.find(hdl);

        if (it == m_connections.end()) {
            throw std::invalid_argument("No data avaliable for session");
        }

        return it->second;
    }
private:
    typedef std::map<connection_hdl, connection_data, std::owner_less<connection_hdl>> con_list;

    int m_next_sessionid;
    server m_server;
    con_list m_connections;
};

int main() {
    broadcast_server server;
    uint16_t port = 1500;

    std::cout << "Server is running on port " << port << "." << std::endl;
    server.run(port);
}