#include "co_gns_server_bridge.h"

#include <steam/steamnetworkingsockets.h>

#include <algorithm>
#include <array>
#include <cstdio>
#include <cstring>
#include <deque>
#include <mutex>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

namespace
{
    constexpr std::size_t kMaximumMessageBytes = 64uz * 1024uz;
    constexpr std::size_t kReceiveBatchSize = 64;

    std::mutex g_runtimeMutex;
    std::size_t g_runtimeReferenceCount = 0;

    class ServerBridge;

    std::mutex g_ownerMutex;
    std::unordered_map<HSteamListenSocket, ServerBridge*> g_listenOwners;
    std::unordered_map<HSteamNetConnection, ServerBridge*> g_connectionOwners;

    void WriteError(char* buffer, std::size_t bufferSize, const std::string& message)
    {
        if (buffer == nullptr || bufferSize == 0) {
            return;
        }
        std::snprintf(buffer, bufferSize, "%s", message.c_str());
    }

    bool AcquireRuntime(std::string& error)
    {
        const std::scoped_lock lock(g_runtimeMutex);
        if (g_runtimeReferenceCount > 0) {
            ++g_runtimeReferenceCount;
            return true;
        }

        SteamNetworkingErrMsg errorMessage{};
        if (!GameNetworkingSockets_Init(nullptr, errorMessage)) {
            error = errorMessage[0] ? std::string{ errorMessage } : "GameNetworkingSockets_Init failed.";
            return false;
        }
        g_runtimeReferenceCount = 1;
        return true;
    }

    void ReleaseRuntime()
    {
        const std::scoped_lock lock(g_runtimeMutex);
        if (g_runtimeReferenceCount == 0) {
            return;
        }
        --g_runtimeReferenceCount;
        if (g_runtimeReferenceCount == 0) {
            GameNetworkingSockets_Kill();
        }
    }

    struct QueuedEvent
    {
        co_gns_event metadata{};
        std::vector<std::uint8_t> payload;
    };

    class ServerBridge
    {
    public:
        ~ServerBridge()
        {
            Stop();
        }

        bool Start(const char* bindHost, std::uint16_t port, std::string& error)
        {
            if (!AcquireRuntime(error)) {
                return false;
            }
            runtimeHeld_ = true;
            networking_ = SteamNetworkingSockets();
            if (networking_ == nullptr) {
                error = "SteamNetworkingSockets returned no server interface.";
                Stop();
                return false;
            }

            pollGroup_ = networking_->CreatePollGroup();
            if (pollGroup_ == k_HSteamNetPollGroup_Invalid) {
                error = "GameNetworkingSockets failed to create a server poll group.";
                Stop();
                return false;
            }

            SteamNetworkingIPAddr address{};
            address.Clear();
            const std::string host = bindHost != nullptr ? std::string{ bindHost } : std::string{};
            if (host.empty() || host == "0.0.0.0") {
                address.SetIPv4(0U, port);
            } else {
                if (!address.ParseString(host.c_str()) || !address.IsIPv4()) {
                    error = "GNS bind address must be a valid IPv4 address.";
                    Stop();
                    return false;
                }
                address.m_port = port;
            }

            SteamNetworkingConfigValue_t option{};
            option.SetPtr(
                k_ESteamNetworkingConfig_Callback_ConnectionStatusChanged,
                reinterpret_cast<void*>(+ConnectionStatusChanged));
            listenSocket_ = networking_->CreateListenSocketIP(address, 1, std::addressof(option));
            if (listenSocket_ == k_HSteamListenSocket_Invalid) {
                error = "GameNetworkingSockets failed to create a listen socket.";
                Stop();
                return false;
            }

            {
                const std::scoped_lock ownerLock(g_ownerMutex);
                g_listenOwners[listenSocket_] = this;
            }

            SteamNetworkingIPAddr actualAddress{};
            if (!networking_->GetListenSocketAddress(listenSocket_, std::addressof(actualAddress))) {
                error = "GameNetworkingSockets could not report the bound listen address.";
                Stop();
                return false;
            }
            localPort_ = actualAddress.m_port;
            return true;
        }

        void Stop()
        {
            ISteamNetworkingSockets* networking = networking_;
            HSteamListenSocket listenSocket = listenSocket_;
            HSteamNetPollGroup pollGroup = pollGroup_;
            std::vector<HSteamNetConnection> connections;
            {
                const std::scoped_lock lock(mutex_);
                connections.assign(connections_.begin(), connections_.end());
                connections_.clear();
                connectedConnections_.clear();
                events_.clear();
                listenSocket_ = k_HSteamListenSocket_Invalid;
                pollGroup_ = k_HSteamNetPollGroup_Invalid;
                networking_ = nullptr;
                localPort_ = 0;
            }

            {
                const std::scoped_lock ownerLock(g_ownerMutex);
                if (listenSocket != k_HSteamListenSocket_Invalid) {
                    g_listenOwners.erase(listenSocket);
                }
                for (const auto connection : connections) {
                    g_connectionOwners.erase(connection);
                }
            }

            if (networking != nullptr) {
                for (const auto connection : connections) {
                    networking->CloseConnection(connection, 0, nullptr, false);
                }
                if (listenSocket != k_HSteamListenSocket_Invalid) {
                    networking->CloseListenSocket(listenSocket);
                }
                if (pollGroup != k_HSteamNetPollGroup_Invalid) {
                    networking->DestroyPollGroup(pollGroup);
                }
            }

            if (runtimeHeld_) {
                runtimeHeld_ = false;
                ReleaseRuntime();
            }
        }

        std::uint16_t LocalPort() const
        {
            const std::scoped_lock lock(mutex_);
            return localPort_;
        }

        std::uint32_t ConnectionCount() const
        {
            const std::scoped_lock lock(mutex_);
            return static_cast<std::uint32_t>(connectedConnections_.size());
        }

        int Poll(co_gns_event& outEvent, void* payloadBuffer, std::uint32_t payloadCapacity)
        {
            ISteamNetworkingSockets* networking = nullptr;
            {
                const std::scoped_lock lock(mutex_);
                networking = networking_;
            }
            if (networking == nullptr) {
                return -1;
            }

            networking->RunCallbacks();

            bool queueIsEmpty = false;
            {
                const std::scoped_lock lock(mutex_);
                queueIsEmpty = events_.empty();
            }
            if (queueIsEmpty) {
                PumpMessages();
            }

            const std::scoped_lock lock(mutex_);
            if (events_.empty()) {
                std::memset(std::addressof(outEvent), 0, sizeof(outEvent));
                return 0;
            }

            const auto& next = events_.front();
            outEvent = next.metadata;
            if (outEvent.type == CO_GNS_EVENT_MESSAGE && next.payload.size() > payloadCapacity) {
                return -2;
            }
            if (outEvent.type == CO_GNS_EVENT_MESSAGE && !next.payload.empty()) {
                if (payloadBuffer == nullptr) {
                    return -2;
                }
                std::memcpy(payloadBuffer, next.payload.data(), next.payload.size());
            }
            events_.pop_front();
            return 1;
        }

        int Send(
            std::uint32_t connectionId,
            const void* payload,
            std::uint32_t payloadSize,
            std::uint32_t delivery)
        {
            if (payload == nullptr || payloadSize == 0) {
                return CO_GNS_SEND_ERROR;
            }
            if (payloadSize > kMaximumMessageBytes) {
                return CO_GNS_SEND_TOO_LARGE;
            }
            if (delivery != CO_GNS_DELIVERY_UNRELIABLE_SEQUENCED && delivery != CO_GNS_DELIVERY_RELIABLE_ORDERED) {
                return CO_GNS_SEND_ERROR;
            }

            const auto connection = static_cast<HSteamNetConnection>(connectionId);
            ISteamNetworkingSockets* networking = nullptr;
            {
                const std::scoped_lock lock(mutex_);
                if (!connectedConnections_.contains(connection)) {
                    return CO_GNS_SEND_NOT_CONNECTED;
                }
                networking = networking_;
            }
            if (networking == nullptr) {
                return CO_GNS_SEND_NOT_CONNECTED;
            }

            const int flags = delivery == CO_GNS_DELIVERY_UNRELIABLE_SEQUENCED ?
                              k_nSteamNetworkingSend_UnreliableNoDelay :
                              k_nSteamNetworkingSend_ReliableNoNagle;
            const auto result = networking->SendMessageToConnection(
                connection,
                payload,
                payloadSize,
                flags,
                nullptr);
            switch (result) {
            case k_EResultOK:
                return CO_GNS_SEND_SENT;
            case k_EResultIgnored:
                return CO_GNS_SEND_DROPPED;
            case k_EResultLimitExceeded:
                return CO_GNS_SEND_BACKPRESSURE;
            case k_EResultNoConnection:
            case k_EResultInvalidState:
                return CO_GNS_SEND_NOT_CONNECTED;
            default:
                return CO_GNS_SEND_ERROR;
            }
        }

        int Disconnect(std::uint32_t connectionId, std::int32_t reason, const char* debug)
        {
            const auto connection = static_cast<HSteamNetConnection>(connectionId);
            ISteamNetworkingSockets* networking = nullptr;
            {
                const std::scoped_lock lock(mutex_);
                if (!connections_.contains(connection)) {
                    return 0;
                }
                connections_.erase(connection);
                connectedConnections_.erase(connection);
                networking = networking_;
            }
            {
                const std::scoped_lock ownerLock(g_ownerMutex);
                g_connectionOwners.erase(connection);
            }
            if (networking == nullptr) {
                return 0;
            }
            return networking->CloseConnection(connection, reason, debug, false) ? 1 : 0;
        }

    private:
        static ServerBridge* FindOwner(const SteamNetConnectionStatusChangedCallback_t& info)
        {
            const std::scoped_lock ownerLock(g_ownerMutex);
            const auto connectionOwner = g_connectionOwners.find(info.m_hConn);
            if (connectionOwner != g_connectionOwners.end()) {
                return connectionOwner->second;
            }
            const auto listenOwner = g_listenOwners.find(info.m_info.m_hListenSocket);
            return listenOwner == g_listenOwners.end() ? nullptr : listenOwner->second;
        }

        static void ConnectionStatusChanged(SteamNetConnectionStatusChangedCallback_t* info)
        {
            if (info == nullptr) {
                return;
            }
            if (auto* owner = FindOwner(*info); owner != nullptr) {
                owner->OnConnectionStatusChanged(*info);
            }
        }

        static co_gns_event MakeEvent(
            std::uint32_t type,
            HSteamNetConnection connection,
            std::int32_t reason,
            std::uint32_t payloadSize,
            const char* debug)
        {
            co_gns_event event{};
            event.type = type;
            event.connection_id = static_cast<std::uint32_t>(connection);
            event.reason = reason;
            event.payload_size = payloadSize;
            if (debug != nullptr && debug[0] != '\0') {
                std::snprintf(event.debug, sizeof(event.debug), "%s", debug);
            }
            return event;
        }

        void OnConnectionStatusChanged(const SteamNetConnectionStatusChangedCallback_t& info)
        {
            if (info.m_info.m_eState == k_ESteamNetworkingConnectionState_Connecting &&
                info.m_info.m_hListenSocket == listenSocket_) {
                if (networking_->AcceptConnection(info.m_hConn) != k_EResultOK ||
                    !networking_->SetConnectionPollGroup(info.m_hConn, pollGroup_)) {
                    networking_->CloseConnection(info.m_hConn, 0, "Could not admit GNS connection", false);
                    return;
                }
                {
                    const std::scoped_lock lock(mutex_);
                    connections_.insert(info.m_hConn);
                }
                {
                    const std::scoped_lock ownerLock(g_ownerMutex);
                    g_connectionOwners[info.m_hConn] = this;
                }
                return;
            }

            if (info.m_info.m_eState == k_ESteamNetworkingConnectionState_Connected) {
                const std::scoped_lock lock(mutex_);
                if (connections_.contains(info.m_hConn) && connectedConnections_.insert(info.m_hConn).second) {
                    events_.push_back(QueuedEvent{
                        MakeEvent(CO_GNS_EVENT_CONNECTED, info.m_hConn, 0, 0, nullptr),
                        {}
                    });
                }
                return;
            }

            if (info.m_info.m_eState != k_ESteamNetworkingConnectionState_ClosedByPeer &&
                info.m_info.m_eState != k_ESteamNetworkingConnectionState_ProblemDetectedLocally) {
                return;
            }

            bool knownConnection = false;
            {
                const std::scoped_lock lock(mutex_);
                knownConnection = connections_.erase(info.m_hConn) > 0;
                connectedConnections_.erase(info.m_hConn);
                if (knownConnection) {
                    events_.push_back(QueuedEvent{
                        MakeEvent(
                            CO_GNS_EVENT_DISCONNECTED,
                            info.m_hConn,
                            info.m_info.m_eEndReason,
                            0,
                            info.m_info.m_szEndDebug),
                        {}
                    });
                }
            }
            if (!knownConnection) {
                return;
            }
            {
                const std::scoped_lock ownerLock(g_ownerMutex);
                g_connectionOwners.erase(info.m_hConn);
            }
            networking_->CloseConnection(info.m_hConn, 0, nullptr, false);
        }

        void PumpMessages()
        {
            ISteamNetworkingSockets* networking = nullptr;
            HSteamNetPollGroup pollGroup = k_HSteamNetPollGroup_Invalid;
            {
                const std::scoped_lock lock(mutex_);
                networking = networking_;
                pollGroup = pollGroup_;
            }
            if (networking == nullptr || pollGroup == k_HSteamNetPollGroup_Invalid) {
                return;
            }

            std::array<SteamNetworkingMessage_t*, kReceiveBatchSize> messages{};
            const auto count = networking->ReceiveMessagesOnPollGroup(
                pollGroup,
                messages.data(),
                static_cast<int>(messages.size()));
            if (count <= 0) {
                return;
            }

            for (int index = 0; index < count; ++index) {
                auto* message = messages[static_cast<std::size_t>(index)];
                if (message == nullptr) {
                    continue;
                }
                const auto size = message->m_cbSize > 0 ? static_cast<std::size_t>(message->m_cbSize) : 0uz;
                const auto connection = message->m_conn;
                QueuedEvent queued{};
                if (size > kMaximumMessageBytes) {
                    queued.metadata = MakeEvent(
                        CO_GNS_EVENT_OVERSIZE_MESSAGE,
                        connection,
                        0,
                        static_cast<std::uint32_t>((std::min)(size, static_cast<std::size_t>(UINT32_MAX))),
                        "Inbound GNS message exceeded 64 KiB");
                } else {
                    queued.metadata = MakeEvent(
                        CO_GNS_EVENT_MESSAGE,
                        connection,
                        0,
                        static_cast<std::uint32_t>(size),
                        nullptr);
                    if (size > 0 && message->m_pData != nullptr) {
                        const auto* begin = static_cast<const std::uint8_t*>(message->m_pData);
                        queued.payload.assign(begin, begin + size);
                    }
                }
                {
                    const std::scoped_lock lock(mutex_);
                    if (connections_.contains(connection)) {
                        events_.push_back(std::move(queued));
                    }
                }
                message->Release();
            }
        }

        mutable std::mutex mutex_;
        ISteamNetworkingSockets* networking_{ nullptr };
        HSteamListenSocket listenSocket_{ k_HSteamListenSocket_Invalid };
        HSteamNetPollGroup pollGroup_{ k_HSteamNetPollGroup_Invalid };
        std::unordered_set<HSteamNetConnection> connections_;
        std::unordered_set<HSteamNetConnection> connectedConnections_;
        std::deque<QueuedEvent> events_;
        std::uint16_t localPort_{ 0 };
        bool runtimeHeld_{ false };
    };
}

extern "C"
{
    int co_gns_server_create(
        const char* bind_host,
        uint16_t port,
        co_gns_server_handle* out_handle,
        char* error_buffer,
        size_t error_buffer_size)
    {
        if (out_handle == nullptr) {
            WriteError(error_buffer, error_buffer_size, "out_handle is required.");
            return 0;
        }
        *out_handle = nullptr;
        auto* server = new ServerBridge();
        std::string error;
        if (!server->Start(bind_host, port, error)) {
            delete server;
            WriteError(error_buffer, error_buffer_size, error);
            return 0;
        }
        *out_handle = server;
        WriteError(error_buffer, error_buffer_size, "");
        return 1;
    }

    void co_gns_server_destroy(co_gns_server_handle handle)
    {
        delete static_cast<ServerBridge*>(handle);
    }

    uint16_t co_gns_server_local_port(co_gns_server_handle handle)
    {
        const auto* server = static_cast<ServerBridge*>(handle);
        return server == nullptr ? 0 : server->LocalPort();
    }

    uint32_t co_gns_server_connection_count(co_gns_server_handle handle)
    {
        const auto* server = static_cast<ServerBridge*>(handle);
        return server == nullptr ? 0 : server->ConnectionCount();
    }

    int co_gns_server_poll(
        co_gns_server_handle handle,
        co_gns_event* out_event,
        void* payload_buffer,
        uint32_t payload_capacity)
    {
        auto* server = static_cast<ServerBridge*>(handle);
        if (server == nullptr || out_event == nullptr) {
            return -1;
        }
        return server->Poll(*out_event, payload_buffer, payload_capacity);
    }

    int co_gns_server_send(
        co_gns_server_handle handle,
        uint32_t connection_id,
        const void* payload,
        uint32_t payload_size,
        uint32_t delivery)
    {
        auto* server = static_cast<ServerBridge*>(handle);
        return server == nullptr ? CO_GNS_SEND_ERROR :
               server->Send(connection_id, payload, payload_size, delivery);
    }

    int co_gns_server_disconnect(
        co_gns_server_handle handle,
        uint32_t connection_id,
        int32_t reason,
        const char* debug)
    {
        auto* server = static_cast<ServerBridge*>(handle);
        return server == nullptr ? 0 : server->Disconnect(connection_id, reason, debug);
    }
}