NAK_CVR_Mods/AvatarScale/Networking/ModNetwork.cs

using ABI_RC.Core.Networking;
using ABI_RC.Systems.ModNetwork;
using DarkRift;
using MelonLoader;
using NAK.AvatarScaleMod.AvatarScaling;
using UnityEngine;

namespace NAK.AvatarScaleMod.Networking;

// overcomplicated, but functional
// a

public static class ModNetwork
{
    #region Constants

    private const string ModId = "MelonMod.NAK.AvatarScaleMod";
    private const float SendRateLimit = 0.25f;
    private const float ReceiveRateLimit = 0.2f;
    private const int MaxWarnings = 2;
    private const float TimeoutDuration = 10f;
    
    private class QueuedMessage
    {
        public MessageType Type { get; set; }
        public float Height { get; set; }
        public string TargetPlayer { get; set; }
        public string Sender { get; set; }
    }

    #endregion

    #region Private State

    private static readonly Dictionary<string, QueuedMessage> OutboundQueue = new();
    private static float LastSentTime;

    private static readonly Dictionary<string, QueuedMessage> InboundQueue = new();
    private static readonly Dictionary<string, float> LastReceivedTimes = new();

    private static readonly Dictionary<string, int> UserWarnings = new();
    private static readonly Dictionary<string, float> UserTimeouts = new();

    #endregion

    #region Enums

    private enum MessageType : byte
    {
        SyncHeight = 0, // just send height
        RequestHeight = 1 // send height, request height back
    }

    #endregion

    #region Mod Network Internals

    internal static void Subscribe()
    {
        ModNetworkManager.Subscribe(ModId, OnMessageReceived);
    }

    internal static void Update()
    {
        ProcessOutboundQueue();
        ProcessInboundQueue();
    }

    private static void SendMessage(MessageType messageType, float height, string playerId = null)
    {
        if (!IsConnectedToGameNetwork())
            return;

        if (!string.IsNullOrEmpty(playerId))
        {
            // to specific user
            using ModNetworkMessage modMsg = new(ModId, playerId);
            modMsg.Write((byte)messageType);
            modMsg.Write(height);
            modMsg.Send();
        }
        else
        {
            // to all users
            using ModNetworkMessage modMsg = new(ModId);
            modMsg.Write((byte)messageType);
            modMsg.Write(height);
            modMsg.Send();
        }

        var typeDesc = messageType == MessageType.SyncHeight ? "height" : "height request";
        MelonLogger.Msg($"Sending {typeDesc}: {height}");
    }

    private static void OnMessageReceived(ModNetworkMessage msg)
    {
        msg.Read(out byte msgTypeRaw);
        msg.Read(out float receivedHeight);

        if (!Enum.IsDefined(typeof(MessageType), msgTypeRaw))
            return;

        // User is in timeout
        if (UserTimeouts.TryGetValue(msg.Sender, out var timeoutEnd) && Time.time < timeoutEnd)
            return;

        if (IsRateLimited(msg.Sender))
            return;

        QueuedMessage inboundMessage = new()
        {
            Type = (MessageType)msgTypeRaw,
            Height = receivedHeight,
            Sender = msg.Sender
        };

        InboundQueue[msg.Sender] = inboundMessage;
        MelonLogger.Msg($"Received message from {msg.Sender}: {receivedHeight}");
    }

    #endregion

    #region Public Methods

    public static void SendNetworkHeight(float newHeight)
    {
        OutboundQueue["global"] = new QueuedMessage { Type = MessageType.SyncHeight, Height = newHeight };
    }

    public static void RequestHeightSync()
    {
        var myCurrentHeight = AvatarScaleManager.Instance.GetHeight();
        if (myCurrentHeight > 0)
            OutboundQueue["global"] = new QueuedMessage { Type = MessageType.RequestHeight, Height = myCurrentHeight };
    }

    #endregion

    #region Outbound Height Queue

    private static void ProcessOutboundQueue()
    {
        if (OutboundQueue.Count == 0 || Time.time - LastSentTime < SendRateLimit)
            return;

        foreach (QueuedMessage message in OutboundQueue.Values)
            SendMessage(message.Type, message.Height, message.TargetPlayer);

        OutboundQueue.Clear();
        LastSentTime = Time.time;
    }

    #endregion

    #region Inbound Height Queue

    private static bool IsRateLimited(string userId)
    {
        // Rate-limit checking
        if (LastReceivedTimes.TryGetValue(userId, out var lastReceivedTime) &&
            Time.time - lastReceivedTime < ReceiveRateLimit)
        {
            if (UserWarnings.TryGetValue(userId, out var warnings))
            {
                warnings++;
                UserWarnings[userId] = warnings;

                if (warnings >= MaxWarnings)
                {
                    UserTimeouts[userId] = Time.time + TimeoutDuration;
                    MelonLogger.Msg($"User is sending height updates too fast! Applying 10s timeout... : {userId}");
                    return true;
                }
            }
            else
            {
                UserWarnings[userId] = 1;
            }

            return true;
        }

        LastReceivedTimes[userId] = Time.time;
        UserWarnings.Remove(userId); // Reset warnings
        // MelonLogger.Msg($"Clearing timeout from user : {userId}");
        return false;
    }

    private static void ProcessInboundQueue()
    {
        foreach (QueuedMessage message in InboundQueue.Values)
            switch (message.Type)
            {
                case MessageType.RequestHeight:
                {
                    var myCurrentHeight = AvatarScaleManager.Instance.GetHeight();
                    if (myCurrentHeight > 0)
                        OutboundQueue[message.Sender] = new QueuedMessage
                        {
                            Type = MessageType.SyncHeight,
                            Height = myCurrentHeight,
                            TargetPlayer = message.Sender
                        };

                    AvatarScaleManager.Instance.OnNetworkHeightUpdateReceived(message.Sender, message.Height);
                    break;
                }
                case MessageType.SyncHeight:
                    AvatarScaleManager.Instance.OnNetworkHeightUpdateReceived(message.Sender, message.Height);
                    break;
                default:
                    AvatarScaleMod.Logger.Error($"Invalid message type received from: {message.Sender}");
                    break;
            }

        InboundQueue.Clear();
    }

    #endregion

    #region Private Methods

    private static bool IsConnectedToGameNetwork()
    {
        return NetworkManager.Instance != null
               && NetworkManager.Instance.GameNetwork != null
               && NetworkManager.Instance.GameNetwork.ConnectionState == ConnectionState.Connected;
    }

    #endregion
}