// The server for Smack a Brother. Players connect over a WebSocket (behind nginx, which handles TLS), make a
// lobby or join one by its code, and the server passes WebRTC signaling between a lobby's members so they can connect
// to each other directly. Nothing of a match passes through here. See src/protocol.h for the messages.
//
// smack-a-brother-server [port]   (listens on 127.0.0.1, port 47800 by default)

#include "protocol.h"
#include <rtc/rtc.hpp>
#include <simdjson.h>

#include <algorithm>
#include <csignal>
#include <cstdio>
#include <map>
#include <mutex>
#include <random>
#include <thread>
#include <vector>

constexpr uint16_t PORT_DEFAULT = 47800;
constexpr size_t MESSAGE_MAX = 16384;
constexpr int CODE_TRIES = 100;

// Someone connected: their id, name, socket, and the lobby they're in (empty until they make or join one).
struct Member
{
    std::string name;
    std::shared_ptr<rtc::WebSocket> socket;
    std::string code;
};

// Every member and lobby (codes to member ids, in the order they joined, the first hosting). The WebSocket callbacks
// run on the library's threads, so everything here is touched under the lock.
static std::mutex lock;
static std::map<int, Member> members;
static std::map<std::string, std::vector<int>> lobbies;
static simdjson::dom::parser parser;
static std::mt19937 rng(std::random_device{}());
static int nextId = 1;

static void send(int id, const std::string& message)
{
    auto member = members.find(id);
    if (member != members.end() && member->second.socket->isOpen())
        member->second.socket->send(message);
}

static void error_send(int id, const char* message)
{
    send(id, "{\"type\":\"error\",\"message\":" + json_quote(message) + "}");
}

// A name as the lobby shows it: printable characters only, cut to length, or "Player" when nothing's left.
static std::string name_clean(std::string_view raw)
{
    std::string name;
    for (char c : raw)
        if (c >= ' ' && c <= '~' && (int)name.size() < PLAYER_NAME_MAX)
            name += c;
    return name.empty() ? "Player" : name;
}

static std::string code_make()
{
    std::uniform_int_distribution<int> letter(0, (int)sizeof(LOBBY_CODE_LETTERS) - 2);
    for (int attempt = 0; attempt < CODE_TRIES; attempt++)
    {
        std::string code;
        for (int i = 0; i < LOBBY_CODE_LENGTH; i++)
            code += LOBBY_CODE_LETTERS[letter(rng)];
        if (!lobbies.contains(code))
            return code;
    }
    return "";
}

// Puts a member into a lobby and tells everyone: the newcomer gets the whole lobby, the rest get the newcomer.
static void lobby_enter(int id, const std::string& code)
{
    std::vector<int>& lobby = lobbies[code];
    lobby.push_back(id);
    members[id].code = code;
    std::string list;
    for (int other : lobby)
        list += (list.empty() ? "" : ",") + std::string("{\"id\":") + std::to_string(other) +
                ",\"name\":" + json_quote(members[other].name) + "}";
    send(id, "{\"type\":\"welcome\",\"id\":" + std::to_string(id) + ",\"code\":" + json_quote(code) +
                 ",\"host\":" + std::to_string(lobby.front()) + ",\"members\":[" + list + "]}");
    for (int other : lobby)
        if (other != id)
            send(other, "{\"type\":\"member_join\",\"id\":" + std::to_string(id) +
                            ",\"name\":" + json_quote(members[id].name) + "}");
}

// Takes a member out of its lobby, telling the rest (and who hosts now, and why if they were kicked). An empty lobby is
// gone.
static void lobby_exit(int id, const std::string& kick = "", bool isKicked = false)
{
    std::string code = members[id].code;
    if (code.empty())
        return;
    members[id].code.clear();
    std::vector<int>& lobby = lobbies[code];
    std::erase(lobby, id);
    if (lobby.empty())
    {
        lobbies.erase(code);
        std::printf("Lobby %s closed (%zu open)\n", code.c_str(), lobbies.size());
        return;
    }
    std::string reason = isKicked ? ",\"kicked\":true,\"reason\":" + json_quote(kick) : "";
    for (int other : lobby)
        send(other, "{\"type\":\"member_leave\",\"id\":" + std::to_string(id) +
                        ",\"host\":" + std::to_string(lobby.front()) + reason + "}");
}

static void message_handle(int id, const std::string& text)
{
    simdjson::dom::element message;
    std::string_view type;
    if (parser.parse(text).get(message) || message["type"].get(type))
        return error_send(id, "Bad message");

    // Answered, so proxies see the connection busy both ways.
    if (type == "ping")
        return send(id, "{\"type\":\"pong\"}");
    if (type == "create" || type == "join")
    {
        int64_t version = 0;
        std::string_view name;
        if (message["version"].get(version) || version != PROTOCOL_VERSION)
            return error_send(id, "Different game version");
        if (!members[id].code.empty())
            return error_send(id, "Already in a lobby");
        members[id].name = name_clean(message["name"].get(name) ? "" : name);
        if (type == "create")
        {
            std::string code = code_make();
            if (code.empty())
                return error_send(id, "No free lobby codes");
            std::printf("Lobby %s opened (%zu open)\n", code.c_str(), lobbies.size() + 1);
            return lobby_enter(id, code);
        }
        std::string_view raw;
        if (message["code"].get(raw))
            return error_send(id, "No code");
        std::string code;
        for (char c : raw)
            code += (char)std::toupper((unsigned char)c);
        auto lobby = lobbies.find(code);
        if (lobby == lobbies.end())
            return error_send(id, "No lobby with that code");
        if ((int)lobby->second.size() >= LOBBY_MEMBERS_MAX)
            return error_send(id, "That lobby is full");
        return lobby_enter(id, code);
    }
    if (type == "leave")
        return lobby_exit(id);
    if (type == "kick")
    {
        // Only a lobby's host can kick, and only someone else in their lobby. The kicked member is told why, then is
        // out (their connection to the server stays; they can join a lobby again).
        int64_t target = 0;
        std::string_view raw;
        std::string& code = members[id].code;
        if (message["id"].get(target) || code.empty() || lobbies[code].front() != id || target == id)
            return error_send(id, "Only the host can kick");
        auto kicked = members.find((int)target);
        if (kicked == members.end() || kicked->second.code != code)
            return error_send(id, "No such player in the lobby");
        std::string reason = text_clean(message["reason"].get(raw) ? "" : raw, KICK_REASON_MAX);
        send((int)target, "{\"type\":\"kicked\",\"reason\":" + json_quote(reason) + "}");
        std::printf("Lobby %s: %d kicked %d\n", code.c_str(), id, (int)target);
        return lobby_exit((int)target, reason, true);
    }
    if (type == "signal")
    {
        // Passed on as it came, only between members of the same lobby.
        int64_t to = 0;
        simdjson::dom::element payload;
        if (message["to"].get(to) || message["payload"].get(payload))
            return error_send(id, "Bad signal");
        auto target = members.find((int)to);
        if (target == members.end() || members[id].code.empty() || target->second.code != members[id].code)
            return;
        send((int)to,
             "{\"type\":\"signal\",\"from\":" + std::to_string(id) + ",\"payload\":" + simdjson::minify(payload) + "}");
        return;
    }
    error_send(id, "Unknown message");
}

static volatile std::sig_atomic_t isStopping = 0;

int main(int argc, char** argv)
{
    std::setvbuf(stdout, nullptr, _IOLBF, 0);
    rtc::InitLogger(rtc::LogLevel::Warning);
    rtc::WebSocketServerConfiguration config;
    config.port = argc > 1 ? (uint16_t)std::atoi(argv[1]) : PORT_DEFAULT;
    config.bindAddress = "127.0.0.1";
    config.maxMessageSize = MESSAGE_MAX;
    rtc::WebSocketServer server(config);

    server.onClient([](std::shared_ptr<rtc::WebSocket> socket) {
        std::lock_guard guard(lock);
        int id = nextId++;
        members[id] = {.socket = socket};
        socket->onMessage([id](rtc::message_variant data) {
            if (!std::holds_alternative<std::string>(data))
                return;
            std::lock_guard guard(lock);
            message_handle(id, std::get<std::string>(data));
        });
        socket->onClosed([id]() {
            std::lock_guard guard(lock);
            if (!members.contains(id))
                return;
            lobby_exit(id);
            members.erase(id);
        });
    });
    std::printf("Server listening on 127.0.0.1:%d\n", server.port());

    std::signal(SIGINT, [](int) { isStopping = 1; });
    std::signal(SIGTERM, [](int) { isStopping = 1; });
    while (!isStopping)
        std::this_thread::sleep_for(std::chrono::milliseconds(200));
    std::printf("Server stopping\n");
    // Closing calls back into onClosed, which takes the lock, so the sockets are gathered first.
    std::vector<std::shared_ptr<rtc::WebSocket>> sockets;
    {
        std::lock_guard guard(lock);
        for (auto& [id, member] : members)
            sockets.push_back(member.socket);
    }
    for (auto& socket : sockets)
        socket->close();
    return 0;
}