BoneSync/BoneSync/Networking/NetworkMessage.cs

using BoneSync.Data;
using BoneSync.Networking.Transport;
using BoneSync.Sync;
using MelonLoader;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Threading.Tasks;
#if TEST
using Xunit;
#endif

namespace BoneSync.Networking
{
    public class PacketTypeAttribute : Attribute
    {
        public PacketTypeAttribute(PacketType type)
        {
            packetType = type;
        }

        public PacketType packetType { get; }
    }

    public class PacketReliabilityAttribute : Attribute
    {
        public PacketReliabilityAttribute(PacketReliability reliability)
        {
            this.reliability = reliability;
        }

        public PacketReliability reliability { get; }
    }

    public class PacketCatchupAttribute : Attribute
    {
        public int order { get; }
        public PacketCatchupAttribute(int order)
        {
            this.order = order;
        }
    }

    public class AlwaysExecuteAttribute : Attribute
    {
        public AlwaysExecuteAttribute() {}
    }

    public abstract class NetworkMessage
    {
        public ulong senderId
        {
            private set;
            get;
        }
        private static bool PacketTypesRegistered => PacketTypeMap.Count > 0;

        internal static Dictionary<PacketType, Type> PacketTypeMap = new Dictionary<PacketType, Type>();

        internal PacketType _packetType;
        internal PacketReliability? _reliability;
        internal int? _catchupOrder;
        internal bool? _alwaysExecute;
        internal ByteEncoder byteEncoder = new ByteEncoder();

        public byte[] GetBytes() => byteEncoder.ToArray();

        public bool GetAlwaysExecute()
        {
            if (_alwaysExecute.HasValue) return _alwaysExecute.Value;
            AlwaysExecuteAttribute alwaysExecuteAttribute = GetType().GetCustomAttribute<AlwaysExecuteAttribute>();
            _alwaysExecute = alwaysExecuteAttribute != null;
            return _alwaysExecute.Value;
        }
        public PacketType GetPacketType()
        {
            RegisterPacketTypes();
            if (_packetType == PacketType.Unknown)
            {
                string typeName = GetType().Name;
                KeyValuePair<PacketType, Type> found = PacketTypeMap.FirstOrDefault(x => x.Value.Name == typeName);
                if (found.Value == null)
                {
                    throw new Exception("Cannot find packetType for " + typeName);
                }
                _packetType = found.Key;
            }
            return _packetType;
        }

        public int GetCatchupOrder()
        {
            if (_catchupOrder.HasValue) return _catchupOrder.Value;
            PacketCatchupAttribute catchupAttribute = GetType().GetCustomAttribute<PacketCatchupAttribute>();
            if (catchupAttribute == null)
            {
                _catchupOrder = -1;
                return -1;
            }

            _catchupOrder = catchupAttribute.order;
            return _catchupOrder.Value;
        }

        public PacketReliability GetPacketReliability()
        {
            if (_reliability.HasValue) return _reliability.Value;
            PacketReliabilityAttribute reliabilityAttribute = GetType().GetCustomAttribute<PacketReliabilityAttribute>();
            if (reliabilityAttribute == null)
            {
                _reliability = PacketReliability.Reliable;
                return PacketReliability.Reliable;
            }

            _reliability = reliabilityAttribute.reliability;
            return _reliability.Value;
        }

        public static void RegisterPacketTypes()
        {
            if (PacketTypesRegistered)
            {
                return;
            }
            Assembly assembly = Assembly.GetExecutingAssembly();
            Type[] types = assembly.GetTypes();
            foreach (Type type in types)
            {

                if (type.IsSubclassOf(typeof(NetworkMessage)))
                {
                    PacketType? packetType = type.GetCustomAttributesData().Where(x => x.AttributeType == typeof(PacketTypeAttribute)).Select(x => (PacketType)x.ConstructorArguments[0].Value).FirstOrDefault();
                    PacketTypeMap[packetType.Value] = type;
                }
            }
        }
        public static NetworkMessage ParsePacket(Packet packet)
        {
            //SyncLogger.Msg("Received packet of type " + packet.Info.packetType + " from " + packet.Info.senderId + " Length: " + packet.Data.Length);
            // find a class that can parse this packet using Reflection
            // and return it
            if (!PacketTypeMap.ContainsKey(packet.Info.packetType))
            {
                throw new Exception("No class found for packet type '" + packet.Info.packetType+"'");
            }
            
            Type type = PacketTypeMap[packet.Info.packetType];
            // get the constructor that takes a Packet
            ConstructorInfo constructor = type.GetConstructor(new Type[] { typeof(Packet) }) ?? throw new Exception("No constructor found for type " + type.Name);
            NetworkMessage networkMessage = (NetworkMessage)constructor.Invoke(new object[] { packet });
            networkMessage.senderId = packet.Info.senderId;
            return networkMessage;
        }

        

        public Packet GetPacket(ulong senderId, ulong receiverId)
        {
            PacketInfo packetInfo = new PacketInfo(senderId, receiverId, GetPacketType(), GetPacketReliability(), 0, (byte)SceneSync.CurrentSceneIndex);
            Packet packet = new Packet(packetInfo, GetBytes());
            return packet;
        }
        public void Broadcast()
        {
            Send(BoneSync.transport.BROADCAST_ID);
        }

        public void SendToHost()
        {
            Send(BoneSync.lobby.GetHostId());
        }
        public void Send(ulong receiverId)
        {
            if (BoneSync.IsConnected == false)
            {
                SyncLogger.Warning("Cannot send packet, not connected to lobby");
                return;
            }
            ulong senderId = BoneSync.lobby.GetLocalId();
            Packet packet = GetPacket(senderId, receiverId);
            BoneSync.transport.Send(packet);
        }

        public virtual void Execute()
        {
            SyncLogger.Warning("Execute not implemented for " + GetType().Name);
        }

    }
    #if TEST
    public class NetworkMessageTests
    {
        [Fact]
        public void TestRegisterPackets()
        {
            NetworkMessage.RegisterPacketTypes();
            Assert.NotEmpty(NetworkMessage.PacketTypeMap);
        }
    }
    #endif
}