#include "net.h"
#include "game.h"
#include <SDL3/SDL.h>
#include <functional>
#include <map>
#include <mutex>
#include <rtc/rtc.hpp>
#include <simdjson.h>
#ifdef SDL_PLATFORM_EMSCRIPTEN
#include <emscripten/eventloop.h>
#endif

constexpr const char STUN_SERVER[] = "stun:stun.l.google.com:19302";
constexpr const char CHANNEL_LABEL[] = "game";
constexpr const char INPUTS_LABEL[] = "inputs";
constexpr Uint64 PING_INTERVAL = 1000;
constexpr Uint64 HEARTBEAT_INTERVAL = 20000;
constexpr float NS_PER_MS = 1e6f;

// The WebRTC connection to the server and, once they're there, its two data channels (see Net).
struct PeerLink
{
    std::shared_ptr<rtc::PeerConnection> connection;
    std::shared_ptr<rtc::DataChannel> channel;
    std::shared_ptr<rtc::DataChannel> inputs;

    bool is_open() const
    {
        return channel && channel->isOpen() && inputs && inputs->isOpen();
    }
};

struct Net::Impl
{
    std::shared_ptr<rtc::WebSocket> socket;
    PeerLink link;
    std::string pending;
    simdjson::dom::parser parser;
    Uint64 nextPing = 0;
    // The network simulator (SAB_NET_LAG milliseconds on everything sent, SAB_NET_LOSS percent of input packets
    // dropped), for trying rollback on one machine: sends wait here until they're due.
    struct Delayed
    {
        Uint64 due;
        bool isReliable;
        std::string bytes;
        bool isText = false;
    };
    std::vector<Delayed> delayed;
    Uint64 lag = 0;
    int loss = 0;
    Uint32 lossSeed = 1;
    Uint64 nextHeartbeat = 0;
    bool isHeartbeating = false;

    // What the network threads report, waiting for the main thread.
    std::mutex lock;
    std::vector<std::function<void()>> posted;

    // On desktop, the network calls back on its own threads, so its news waits for the main thread's next update. The
    // browser calls back on the main thread already, and keeps doing so in a hidden tab, where frames (and so updates)
    // stop; there it's handled at once, so a hidden tab still answers a connection being made to it.
    void post(std::function<void()> work)
    {
#ifdef SDL_PLATFORM_EMSCRIPTEN
        work();
#else
        std::lock_guard guard(lock);
        posted.push_back(std::move(work));
#endif
    }
};

Net::Net() : impl(std::make_unique<Impl>())
{
    if (const char* lag = SDL_getenv("SAB_NET_LAG"))
        impl->lag = SDL_atoi(lag);
    if (const char* loss = SDL_getenv("SAB_NET_LOSS"))
        impl->loss = SDL_atoi(loss);
#ifndef SDL_PLATFORM_EMSCRIPTEN
    rtc::InitLogger(rtc::LogLevel::Warning);
#endif
}

// Pings the server so proxies between (which close connections that go quiet) keep it open. The server answers,
// so it's busy both ways.
void Net::heartbeat()
{
    if (impl->socket && impl->socket->isOpen())
        impl->socket->send(std::string("{\"type\":\"ping\"}"));
}

Net::~Net()
{
    shutdown();
}

// Closes every connection, waiting for nothing; anything still queued is dropped.
void Net::shutdown()
{
    if (impl->link.connection)
        impl->link.connection->close();
    impl->link = {};
    if (impl->socket)
        impl->socket->close();
    impl->socket.reset();
    std::lock_guard guard(impl->lock);
    impl->posted.clear();
}

void Net::lobby_create(Game& game)
{
    request_send(game, "{\"type\":\"create\",\"name\":" + json_quote(game.settings.name) +
                           ",\"version\":" + std::to_string(PROTOCOL_VERSION) + "}");
}

void Net::lobby_join(Game& game, const char* lobby)
{
    request_send(game, "{\"type\":\"join\",\"code\":" + json_quote(lobby) + ",\"name\":" +
                           json_quote(game.settings.name) + ",\"version\":" + std::to_string(PROTOCOL_VERSION) + "}");
}

// Leaves the lobby, closing the WebRTC connection but keeping the socket to the server.
void Net::lobby_leave()
{
    if (impl->link.connection)
        impl->link.connection->close();
    impl->link = {};
    if (impl->socket && impl->socket->isOpen() && state == NET_LOBBY)
        impl->socket->send(std::string("{\"type\":\"leave\"}"));
    members.clear();
    packets.clear();
    kicks.clear();
    code.clear();
    selfId = hostId = -1;
    state = NET_OFFLINE;
}

// Sends a create or join, connecting to the server first if need be (it goes out once the socket opens). The
// server is the one in Settings (Host), or SAB_SERVER_URL when set (for testing against a local one).
void Net::request_send(Game& game, const std::string& request)
{
    error.clear();
    if (impl->socket && impl->socket->isOpen())
    {
        impl->socket->send(request);
        state = NET_JOINING;
        return;
    }
    impl->pending = request;
    state = NET_CONNECTING;
    if (impl->socket)
        return;
    impl->socket = std::make_shared<rtc::WebSocket>();
#ifdef SDL_PLATFORM_EMSCRIPTEN
    // A browser timer keeps the server connection alive even while the tab is hidden and frames stop.
    if (!impl->isHeartbeating)
        emscripten_set_interval([](void* net) { static_cast<Net*>(net)->heartbeat(); }, HEARTBEAT_INTERVAL, this);
    impl->isHeartbeating = true;
#endif
    Impl* shared = impl.get();
    impl->socket->onOpen([this, shared]() {
        shared->post([this]() {
            if (impl->pending.empty())
                return;
            impl->socket->send(impl->pending);
            impl->pending.clear();
            state = NET_JOINING;
        });
    });
    impl->socket->onMessage([this, shared](rtc::message_variant data) {
        if (std::holds_alternative<std::string>(data))
        {
            shared->post([this, text = std::get<std::string>(data)]() { message_handle(text); });
            return;
        }
        rtc::binary& bytes = std::get<rtc::binary>(data);
        std::string packet(reinterpret_cast<const char*>(bytes.data()), bytes.size());
        shared->post([this, packet]() { packets.push_back(packet); });
    });
    impl->socket->onClosed([this, shared]() {
        shared->post([this]() {
            lobby_leave();
            impl->socket.reset();
            impl->pending.clear();
            if (error.empty())
                error = "Lost the connection to the server";
        });
    });
    impl->socket->onError([this, shared](std::string message) {
        shared->post([this, message]() { error = "Couldn't reach the server (" + message + ")"; });
    });
    const char* url = SDL_getenv("SAB_SERVER_URL");
    impl->socket->open(url ? url : game.settings.serverUrl);
}

NetMember* Net::member_get(int id)
{
    for (NetMember& member : members)
        if (member.id == id)
            return &member;
    return nullptr;
}

// One message from the server.
void Net::message_handle(const std::string& text)
{
    simdjson::dom::element message;
    std::string_view type;
    if (impl->parser.parse(text).get(message) || message["type"].get(type))
        return;
    int64_t number = 0;
    std::string_view field;

    if (type == "error")
    {
        error = message["message"].get(field) ? "Server error" : std::string(field);
        if (state != NET_LOBBY)
            state = NET_OFFLINE;
        return;
    }
    if (type == "welcome")
    {
        members.clear();
        selfId = message["id"].get(number) ? -1 : (int)number;
        hostId = message["host"].get(number) ? -1 : (int)number;
        code = message["code"].get(field) ? "" : std::string(field);
        simdjson::dom::array list;
        if (!message["members"].get(list))
            for (simdjson::dom::element entry : list)
            {
                NetMember member;
                member.id = entry["id"].get(number) ? -1 : (int)number;
                member.name = entry["name"].get(field) ? "" : std::string(field);
                members.push_back(member);
            }
        state = NET_LOBBY;
        link_open();
        return;
    }
    if (type == "member_join")
    {
        NetMember member;
        member.id = message["id"].get(number) ? -1 : (int)number;
        member.name = message["name"].get(field) ? "" : std::string(field);
        members.push_back(member);
        return;
    }
    if (type == "kicked")
    {
        // Kicked out by the host: out of the lobby, saying why, here and in the chat.
        std::string reason = text_clean(message["reason"].get(field) ? "" : field, KICK_REASON_MAX);
        std::string text = "You were kicked by the host." + (reason.empty() ? "" : " Reason: " + reason);
        state = NET_OFFLINE;
        lobby_leave();
        error = text;
        notices.push_back(text);
        return;
    }
    if (type == "member_leave")
    {
        int id = message["id"].get(number) ? -1 : (int)number;
        bool isKicked = false;
        if (!message["kicked"].get(isKicked) && isKicked)
            kicks[id] = text_clean(message["reason"].get(field) ? "" : field, KICK_REASON_MAX);
        hostId = message["host"].get(number) ? hostId : (int)number;
        std::erase_if(members, [&](NetMember& member) { return member.id == id; });
        return;
    }
    if (type == "signal")
    {
        simdjson::dom::element payload;
        if (!message["payload"].get(payload))
            signal_handle(simdjson::minify(payload));
    }
}

// Offers the server a WebRTC connection with both data channels, which starts the negotiation. Descriptions and
// candidates go back and forth over the socket. Pongs on the reliable channel time the round trip.
void Net::link_open()
{
    rtc::Configuration config;
    config.iceServers.emplace_back(STUN_SERVER);
    PeerLink& link = impl->link;
    link.connection = std::make_shared<rtc::PeerConnection>(config);
    Impl* shared = impl.get();
    auto signal_send = [this](const std::string& payload) {
        if (impl->socket && impl->socket->isOpen())
            impl->socket->send("{\"type\":\"signal\",\"payload\":" + payload + "}");
    };
    link.connection->onLocalDescription([shared, signal_send](rtc::Description description) {
        std::string payload = "{\"sdp\":" + json_quote(std::string(description)) +
                              ",\"kind\":" + json_quote(description.typeString()) + "}";
        shared->post([signal_send, payload]() { signal_send(payload); });
    });
    link.connection->onLocalCandidate([shared, signal_send](rtc::Candidate candidate) {
        std::string payload =
            "{\"candidate\":" + json_quote(candidate.candidate()) + ",\"mid\":" + json_quote(candidate.mid()) + "}";
        shared->post([signal_send, payload]() { signal_send(payload); });
    });

    rtc::DataChannelInit lossy;
    lossy.reliability.unordered = true;
    lossy.reliability.maxRetransmits = 0;
    link.channel = link.connection->createDataChannel(CHANNEL_LABEL);
    link.inputs = link.connection->createDataChannel(INPUTS_LABEL, lossy);
    for (auto& channel : {link.channel, link.inputs})
        channel->onMessage([this, shared](rtc::message_variant data) {
            if (std::holds_alternative<rtc::binary>(data))
            {
                rtc::binary& bytes = std::get<rtc::binary>(data);
                std::string packet(reinterpret_cast<const char*>(bytes.data()), bytes.size());
                shared->post([this, packet]() { packets.push_back(packet); });
                return;
            }
            std::string& text = std::get<std::string>(data);
            if (!text.starts_with("pong:"))
                return;
            float ping = (SDL_GetTicksNS() - SDL_strtoull(text.c_str() + 5, nullptr, 10)) / NS_PER_MS;
            shared->post([this, ping]() {
                if (NetMember* member = member_get(selfId))
                    member->ping = ping;
            });
        });
}

// A description or candidate from the server for this screen's connection.
void Net::signal_handle(const std::string& text)
{
    simdjson::dom::element payload;
    std::string_view sdp;
    std::string_view kind;
    std::string_view candidate;
    std::string_view mid;
    PeerLink& link = impl->link;
    if (impl->parser.parse(text).get(payload) || !link.connection)
        return;
    if (!payload["sdp"].get(sdp) && !payload["kind"].get(kind))
        link.connection->setRemoteDescription(rtc::Description(std::string(sdp), std::string(kind)));
    else if (!payload["candidate"].get(candidate) && !payload["mid"].get(mid))
        link.connection->addRemoteCandidate(rtc::Candidate(std::string(candidate), std::string(mid)));
}

// Sends the server a game message, held back or dropped first by the network simulator if it's on.
void Net::send(const std::string& bytes, bool isReliable)
{
    if (impl->lag > 0 || impl->loss > 0)
    {
        impl->lossSeed = impl->lossSeed * 1103515245 + 12345;
        if (!isReliable && (int)(impl->lossSeed >> 16) % 100 < impl->loss)
            return;
        if (impl->lag > 0)
        {
            impl->delayed.push_back({SDL_GetTicks() + impl->lag, isReliable, bytes});
            return;
        }
    }
    send_now(bytes, isReliable);
}

// Over the data channels once they're open, or the socket until then.
void Net::send_now(const std::string& bytes, bool isReliable)
{
    auto data = reinterpret_cast<const rtc::byte*>(bytes.data());
    if (is_linked())
        (isReliable ? impl->link.channel : impl->link.inputs)->send(data, bytes.size());
    else if (impl->socket && impl->socket->isOpen() && state == NET_LOBBY)
        impl->socket->send(rtc::binary(data, data + bytes.size()));
}

// Times the round trip to the server over the reliable channel (the network simulator holds it back too).
void Net::ping_send()
{
    std::string text = "ping:" + std::to_string(SDL_GetTicksNS());
    if (impl->lag > 0)
        impl->delayed.push_back({SDL_GetTicks() + impl->lag, true, text, true});
    else if (is_linked())
        impl->link.channel->send(text);
}

// Asks the server to put a member out of the lobby (it only does so for the host).
void Net::kick_send(int member, const std::string& reason)
{
    if (impl->socket && impl->socket->isOpen() && state == NET_LOBBY)
        impl->socket->send("{\"type\":\"kick\",\"id\":" + std::to_string(member) + ",\"reason\":" + json_quote(reason) +
                           "}");
}

// Whether the WebRTC connection to the server is up, with both its channels open.
bool Net::is_linked()
{
    return impl->link.is_open();
}

// Once a frame: handles what the network reported (on desktop), pings the server over the link, and on desktop keeps
// the server connection alive.
void Net::update()
{
    std::vector<std::function<void()>> work;
    {
        std::lock_guard guard(impl->lock);
        work.swap(impl->posted);
    }
    for (auto& job : work)
        job();

    Uint64 now = SDL_GetTicks();
    while (!impl->delayed.empty() && impl->delayed.front().due <= now)
    {
        Impl::Delayed due = impl->delayed.front();
        impl->delayed.erase(impl->delayed.begin());
        if (due.isText && is_linked())
            impl->link.channel->send(due.bytes);
        else if (!due.isText)
            send_now(due.bytes, due.isReliable);
    }
    if (state == NET_LOBBY && now >= impl->nextPing)
    {
        impl->nextPing = now + PING_INTERVAL;
        ping_send();
    }
#ifndef SDL_PLATFORM_EMSCRIPTEN
    if (now >= impl->nextHeartbeat)
    {
        impl->nextHeartbeat = now + HEARTBEAT_INTERVAL;
        heartbeat();
    }
#endif
}