Syntriax.Engine/Engine.Systems/Network/NetworkManager.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;

using Engine.Core;

namespace Engine.Systems.Network;

/// <summary>
/// Intermediary manager that looks up in it's hierarchy for a <see cref="INetworkCommunicator"/> to route/broadcast it's received packets to their destinations.
/// </summary>
/// TODO: I need to peer check this class, I don't exactly remember the state I was in when I was originally writing it and left it uncommented and the current comments are added later on.
/// It's a fairly complex manager that relies heavily on Reflection and lots of generic method delegation which is making it very hard to read back. 
public class NetworkManager : Behaviour, IEnterUniverse, IExitUniverse, INetworkManager
{
    #region Packet Router/Broadcaster to Listener Delegates

    /// <summary>
    /// Behaviour Type → Packet Type → List of <see cref="IPacketListenerClient{T}"/> listener methods (broadcast packets, client-side)
    /// </summary>
    private readonly Dictionary<Type, Dictionary<Type, List<MethodInfo>>> clientBroadcastPacketListenerMethods = [];

    /// <summary>
    /// Behaviour Type → Packet Type → List of <see cref="IPacketListenerServer{T}"/> listener methods (broadcast packets, server-side)
    /// </summary>
    private readonly Dictionary<Type, Dictionary<Type, List<MethodInfo>>> serverBroadcastPacketListenerMethods = [];

    /// <summary>
    /// Behaviour Type → Packet Type → List of <see cref="IPacketListenerClientEntity{T}"/> listener methods (entity packets, client-side)
    /// </summary>
    private readonly Dictionary<Type, Dictionary<Type, List<MethodInfo>>> clientEntityPacketListenerMethods = [];

    /// <summary>
    /// Behaviour Type → Packet Type → List of <see cref="IPacketListenerServerEntity{T}"/> listener methods (entity packets, server-side)
    /// </summary>
    private readonly Dictionary<Type, Dictionary<Type, List<MethodInfo>>> serverEntityPacketListenerMethods = [];

    #endregion 

    #region Packet Router/Broadcaster Events

    /// <summary>
    /// Packet Type → Behaviour.Id → Broadcaster Event (broadcast, client-side)
    /// </summary>
    private readonly Dictionary<Type, Dictionary<string, object>> clientPacketBroadcastEvents = [];

    /// <summary>
    /// Packet Type → Behaviour.Id → Broadcaster Event (broadcast, server-side)
    /// </summary>
    private readonly Dictionary<Type, Dictionary<string, object>> serverPacketBroadcastEvents = [];

    /// <summary>
    /// Maps an <see cref="IEntityNetworkPacket"/> type to a set of routing events,
    /// keyed by <see cref="IIdentifiable.Id"/>, for CLIENT entity listeners.
    /// The packet is routed to the correct <see cref="INetworkEntity"/> instance
    /// by matching <see cref="IEntityNetworkPacket.EntityId"/>.
    /// </summary>
    private readonly Dictionary<Type, Dictionary<string, object>> clientEntityPacketRouterEvents = [];

    /// <summary>
    /// Maps an <see cref="IEntityNetworkPacket"/> type to a set of routing events,
    /// keyed by <see cref="IIdentifiable.Id"/>, for SERVER entity listeners.
    /// The packet is routed to the correct <see cref="INetworkEntity"/> instance
    /// by matching <see cref="IEntityNetworkPacket.EntityId"/>.
    /// </summary>
    private readonly Dictionary<Type, Dictionary<string, object>> serverEntityPacketRouterEvents = [];

    #endregion 

    #region Packet Retrieval Delegates

    /// <summary>
    /// Stores delegates that connect incoming broadcast packets from <see cref="INetworkCommunicator"/>
    /// to <see cref="OnPacketReceived{T}(IConnection, T)"/>. 
    /// These are used to subscribe/unsubscribe from <see cref="INetworkCommunicator"/> events.
    /// </summary>
    private readonly List<PacketRetrievalDelegatePair> broadcastPacketRetrievalSubscriptionDelegates = [];

    /// <summary>
    /// Stores delegates that connect incoming entity packets from <see cref="INetworkCommunicator"/>
    /// to <see cref="OnPacketReceived{T}(IConnection, T)"/>. 
    /// These are used to subscribe/unsubscribe from <see cref="INetworkCommunicator"/> events.
    /// </summary>
    private readonly List<PacketRetrievalDelegatePair> entityPacketRetrievalSubscriptionDelegates = [];

    /// <summary>
    /// Stores delegates that connect incoming all packets from <see cref="INetworkCommunicator"/> to 
    /// <see cref="OnPacketReceived{T}(IConnection, T)"/>. This is a combination of all subscription 
    /// delegates filtered so there are no duplicates packet entries.
    /// </summary>
    private readonly List<PacketRetrievalDelegatePair> uniqueRetrievalSubscriptionDelegates = [];

    #endregion

    #region Method Caches

    /// <summary>
    /// Packet type → <see cref="ClearRouter{T}(object)"/> method.
    /// </summary>
    private readonly Dictionary<Type, MethodInfo> clearRoutesMethods = [];

    /// <summary>
    /// Packet type → <see cref="RegisterBroadcastPacketListenerEvent{T}(INetworkEntity, Event{IConnection, T}, NetworkType)"/> method.
    /// </summary>
    private readonly Dictionary<Type, MethodInfo> registerBroadcastPacketListenersMethods = [];

    /// <summary>
    /// Packet type → <see cref="RegisterEntityPacketListenerEvent{T}(INetworkEntity, Event{IConnection, T}, NetworkType)"/> method.
    /// </summary>
    private readonly Dictionary<Type, MethodInfo> registerEntityPacketListenersMethods = [];

    #endregion 

    #region Network Entity Collector

    /// <summary>
    /// All active network <see cref="INetworkEntity"/>, keyed by <see cref="IIdentifiable.Id"/>.
    /// </summary>
    private readonly Dictionary<string, INetworkEntity> _networkEntities = [];
    public IReadOnlyDictionary<string, INetworkEntity> NetworkEntities => _networkEntities;

    /// <summary>
    /// Collector responsible for detecting <see cref="INetworkEntity"/>s entering/leaving the universe.
    /// </summary>
    private readonly BehaviourCollector<INetworkEntity> _networkEntityCollector = new();

    public IBehaviourCollector<INetworkEntity> NetworkEntityCollector => _networkEntityCollector;

    #endregion

    #region Network Communicator

    public INetworkCommunicator NetworkCommunicator
    {
        get;
        set
        {
            if (field == value)
                return;

            INetworkCommunicator? previousCommunicator = field;
            field = value;

            // Unsubscribe packet delegates from old communicator
            if (previousCommunicator is not null)
                InvokeCommunicatorMethods(nameof(INetworkCommunicator.UnsubscribeFromPackets), previousCommunicator);

            // Subscribe packet delegates to new communicator
            if (field is not null)
                InvokeCommunicatorMethods(nameof(INetworkCommunicator.SubscribeToPackets), field);
        }
    } = null!;

    /// <summary>
    /// Dynamically invokes <see cref="INetworkCommunicator.SubscribeToPackets{T}(Event{IConnection, T}.EventHandler)"/> 
    /// or <see cref="INetworkCommunicator.UnsubscribeFromPackets{T}(Event{IConnection, T}.EventHandler)"/>
    /// on the provided <see cref="INetworkCommunicator"/> for all known packet types.
    /// </summary>
    private void InvokeCommunicatorMethods(string methodName, INetworkCommunicator networkCommunicator)
    {
        MethodInfo unsubscribeFromPacketsMethod = typeof(INetworkCommunicator)
            .GetMethod(methodName, BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance)!;

        foreach ((Type packetType, Delegate @delegate) in uniqueRetrievalSubscriptionDelegates)
        {
            MethodInfo genericMethod = unsubscribeFromPacketsMethod.MakeGenericMethod(packetType);
            genericMethod.Invoke(networkCommunicator, [@delegate]);
        }
    }

    #endregion

    ////////////////////////////////////////////////////////////////

    #region Packet Routing/Broadcasting

    /// <summary>
    /// Entry point for ALL incoming packets from the <see cref="NetworkCommunicator"/>.
    /// </summary>
    private void OnPacketReceived<T>(IConnection sender, T entityDataPacket)
    {
        BroadcastPacket(sender, entityDataPacket);

        if (entityDataPacket is IEntityNetworkPacket entityPacket)
            RoutePacket(sender, entityDataPacket, entityPacket);
    }

    private void RoutePacket<T>(IConnection sender, T entityDataPacket, IEntityNetworkPacket entityPacket)
    {
        if (NetworkCommunicator is INetworkCommunicatorClient)
            RoutePacket(clientEntityPacketRouterEvents, entityPacket.EntityId, sender, entityDataPacket);

        if (NetworkCommunicator is INetworkCommunicatorServer)
            RoutePacket(serverEntityPacketRouterEvents, entityPacket.EntityId, sender, entityDataPacket);
    }

    private void BroadcastPacket<T>(IConnection sender, T entityDataPacket)
    {
        if (NetworkCommunicator is INetworkCommunicatorClient)
            BroadcastPacket(clientPacketBroadcastEvents, sender, entityDataPacket);

        if (NetworkCommunicator is INetworkCommunicatorServer)
            BroadcastPacket(serverPacketBroadcastEvents, sender, entityDataPacket);
    }

    private void BroadcastPacket<T>(
        Dictionary<Type, Dictionary<string, object>> packetBroadcasters,
        IConnection sender,
        T entityDataPacket)
    {
        if (!packetBroadcasters.TryGetValue(entityDataPacket!.GetType(), out Dictionary<string, object>? routers))
            return;

        foreach ((string behaviourId, object routerEventReference) in routers)
        {
            Event<IConnection, T> routerEvent = (Event<IConnection, T>)routerEventReference;
            routerEvent.Invoke(sender, entityDataPacket!);
        }
    }

    private void RoutePacket<T>(
        Dictionary<Type, Dictionary<string, object>> packetRouters,
        string entityId,
        IConnection sender,
        T entityDataPacket)
    {
        if (!packetRouters.TryGetValue(entityDataPacket!.GetType(), out Dictionary<string, object>? routers))
            return;

        if (!routers.TryGetValue(entityId, out object? routerEventReference))
            return;

        Event<IConnection, T> routerEvent = (Event<IConnection, T>)routerEventReference;
        routerEvent.Invoke(sender, entityDataPacket!);
    }

    #endregion

    #region Packet Routers

    /// <summary>
    /// Registers routing events for the behaviour based on cached packet listener methods.
    /// </summary>
    private void RegisterPacketRoutersFor(
        INetworkEntity behaviour,
        Dictionary<Type, Dictionary<string, object>> packetRouters,
        Dictionary<Type, Dictionary<Type, List<MethodInfo>>> packetListenerMethods,
        NetworkType networkType,
        Dictionary<Type, MethodInfo> registerPacketListenerListenersMethods)
    {
        if (!packetListenerMethods.TryGetValue(behaviour.GetType(), out Dictionary<Type, List<MethodInfo>>? listenerMethods))
            return;

        foreach (Type packetType in listenerMethods.Keys)
        {
            if (!packetRouters.TryGetValue(packetType, out Dictionary<string, object>? routers))
            {
                routers = [];
                packetRouters.Add(packetType, routers);
            }

            object packetListenerEvent = CreateEventAndRegister(packetType, behaviour, networkType, registerPacketListenerListenersMethods);

            routers.Add(behaviour.Id, packetListenerEvent);
        }
    }

    /// <summary>
    /// Creates an Event<IConnection, TPacket> and attaches listener callbacks.
    /// </summary>
    private object CreateEventAndRegister(
        Type packetType,
        INetworkEntity behaviour,
        NetworkType networkType,
        Dictionary<Type, MethodInfo> registerPacketListenersMethods)
    {
        Type genericEventType = typeof(Event<,>).MakeGenericType(typeof(IConnection), packetType);
        object packetListenerEvent = Activator.CreateInstance(genericEventType)!;

        if (!registerPacketListenersMethods.TryGetValue(packetType, out MethodInfo? registerPacketListenerMethod))
            throw new($"Packet Listener Events for {packetType.Name} has not been cached.");

        registerPacketListenerMethod.Invoke(this, [behaviour, packetListenerEvent, networkType]);
        return packetListenerEvent;
    }

    /// <summary>
    /// Registers broadcast packet listeners on the behaviour.
    /// </summary>
    private static void RegisterBroadcastPacketListenerEvent<T>(
        INetworkEntity behaviour,
        Event<IConnection, T> packetListenerEvent,
        NetworkType networkType)
    {
        switch (networkType)
        {
            case NetworkType.Client:
                packetListenerEvent.AddListener(
                    (sender, packet) => ((IPacketListenerClient<T>)behaviour).OnClientPacketArrived(sender, packet));
                break;

            case NetworkType.Server:
                packetListenerEvent.AddListener(
                    (sender, packet) => ((IPacketListenerServer<T>)behaviour).OnServerPacketArrived(sender, packet));
                break;
        }
    }

    /// <summary>
    /// Registers entity-specific packet listeners on a network behaviour.
    /// </summary>
    private static void RegisterEntityPacketListenerEvent<T>(
        INetworkEntity behaviour,
        Event<IConnection, T> packetListenerEvent,
        NetworkType networkType
    ) where T : IEntityNetworkPacket
    {
        switch (networkType)
        {
            case NetworkType.Client:
                packetListenerEvent.AddListener(
                    (sender, packet) => ((IPacketListenerClientEntity<T>)behaviour).OnEntityClientPacketArrived(sender, packet));
                break;

            case NetworkType.Server:
                packetListenerEvent.AddListener(
                    (sender, packet) => ((IPacketListenerServerEntity<T>)behaviour).OnEntityServerPacketArrived(sender, packet));
                break;
        }
    }

    /// <summary>
    /// Unregisters all routing events associated with the behaviour.
    /// </summary>
    private void UnregisterPacketRoutersFor(
        INetworkEntity behaviour,
        Dictionary<Type, Dictionary<string, object>> packetRouters,
        Dictionary<Type, Dictionary<Type, List<MethodInfo>>> packetListenerMethods)
    {
        if (!packetListenerMethods.TryGetValue(behaviour.GetType(), out Dictionary<Type, List<MethodInfo>>? listenerMethods))
            return;

        foreach ((Type packetType, List<MethodInfo> methods) in listenerMethods)
        {
            if (!packetRouters.TryGetValue(packetType, out Dictionary<string, object>? routers))
                continue;

            if (!routers.TryGetValue(behaviour.Id, out object? routerEventReference))
                continue;

            if (!clearRoutesMethods.TryGetValue(packetType, out MethodInfo? clearRouterMethod))
                continue;

            clearRouterMethod.Invoke(this, [routerEventReference]);
        }
    }

    /// <summary>
    /// Clears all listeners from a router event.
    /// </summary>
    private static void ClearRouter<T>(object routerEventReference)
    {
        Event<IConnection, T> routerEvent = (Event<IConnection, T>)routerEventReference;
        routerEvent.Clear();
    }

    #endregion

    #region Engine Callbacks

    /// <summary>
    /// Called when an <see cref="INetworkEntity"/> enters the universe.
    /// Registers all packet routing for that entity.
    /// </summary>
    private void OnCollected(
        IBehaviourCollector<INetworkEntity> sender,
        IBehaviourCollector<INetworkEntity>.BehaviourCollectedArguments args)
    {
        INetworkEntity collectedBehaviour = args.BehaviourCollected;

        if (!_networkEntities.TryAdd(collectedBehaviour.Id, collectedBehaviour))
            throw new($"Unable to add {collectedBehaviour.Id} to {nameof(NetworkManager)}");

        RegisterPacketRoutersFor(collectedBehaviour, clientPacketBroadcastEvents, clientBroadcastPacketListenerMethods, NetworkType.Client, registerBroadcastPacketListenersMethods);
        RegisterPacketRoutersFor(collectedBehaviour, clientEntityPacketRouterEvents, clientEntityPacketListenerMethods, NetworkType.Client, registerEntityPacketListenersMethods);

        RegisterPacketRoutersFor(collectedBehaviour, serverPacketBroadcastEvents, serverBroadcastPacketListenerMethods, NetworkType.Server, registerBroadcastPacketListenersMethods);
        RegisterPacketRoutersFor(collectedBehaviour, serverEntityPacketRouterEvents, serverEntityPacketListenerMethods, NetworkType.Server, registerEntityPacketListenersMethods);
    }

    /// <summary>
    /// Called when an <see cref="INetworkEntity"/> is removed from the universe.
    /// Cleans up all routing.
    /// </summary>
    private void OnRemoved(IBehaviourCollector<INetworkEntity> sender, IBehaviourCollector<INetworkEntity>.BehaviourRemovedArguments args)
    {
        INetworkEntity removedBehaviour = args.BehaviourRemoved;

        if (!_networkEntities.Remove(removedBehaviour.Id))
            return;

        UnregisterPacketRoutersFor(removedBehaviour, clientPacketBroadcastEvents, clientBroadcastPacketListenerMethods);
        UnregisterPacketRoutersFor(removedBehaviour, clientEntityPacketRouterEvents, clientEntityPacketListenerMethods);

        UnregisterPacketRoutersFor(removedBehaviour, serverPacketBroadcastEvents, serverBroadcastPacketListenerMethods);
        UnregisterPacketRoutersFor(removedBehaviour, serverEntityPacketRouterEvents, serverEntityPacketListenerMethods);
    }

    public void ExitUniverse(IUniverse universe) => _networkEntityCollector.Unassign();

    public void EnterUniverse(IUniverse universe)
    {
        _networkEntityCollector.Assign(universe);
        NetworkCommunicator = BehaviourController.GetRequiredBehaviourInParent<INetworkCommunicator>();
    }

    #endregion

    #region Initialization

    public NetworkManager()
    {
        CacheRetrievalSubscriptionDelegates();
        CacheRegistrationMethods();
        CachePacketListenerMethods();

        _networkEntityCollector.OnCollected.AddListener(OnCollected);
        _networkEntityCollector.OnRemoved.AddListener(OnRemoved);
    }

    /// <summary>
    /// Caches all retrieval subscription delegates for packets.
    /// </summary>
    private void CacheRetrievalSubscriptionDelegates()
    {
        CachePacketRetrievalDelegates(typeof(INetworkPacket), broadcastPacketRetrievalSubscriptionDelegates);
        CachePacketRetrievalDelegates(typeof(IEntityNetworkPacket), entityPacketRetrievalSubscriptionDelegates);

        uniqueRetrievalSubscriptionDelegates.AddRange(broadcastPacketRetrievalSubscriptionDelegates.Concat(entityPacketRetrievalSubscriptionDelegates).DistinctBy(pair => pair.PacketType));
    }

    /// <summary>
    /// Creates delegates for all concrete packet types that forward packets to <see cref="OnPacketReceived{T}(IConnection, T)"/>.
    /// </summary>
    private void CachePacketRetrievalDelegates(Type packetType, List<PacketRetrievalDelegatePair> retrievalDelegates)
    {
        IEnumerable<Type> packetTypes =
            AppDomain.CurrentDomain.GetAssemblies()
                .SelectMany(a => a.GetTypes())
                .Where(t =>
                    packetType.IsAssignableFrom(t) &&
                    !t.IsInterface &&
                    !t.IsAbstract &&
                    !t.IsGenericType
                );

        MethodInfo onPacketArrivedMethod = GetType().GetMethod(nameof(OnPacketReceived), BindingFlags.NonPublic | BindingFlags.Instance)!;

        foreach (Type type in packetTypes)
        {
            MethodInfo genericOnPacketArrivedMethod = onPacketArrivedMethod.MakeGenericMethod(type);

            Type genericDelegateType = typeof(Event<,>.EventHandler).MakeGenericType(typeof(IConnection), type);

            Delegate genericPacketReceivedDelegate = Delegate.CreateDelegate(genericDelegateType, this, genericOnPacketArrivedMethod);

            retrievalDelegates.Add((type, genericPacketReceivedDelegate));
        }
    }

    /// <summary>
    /// Caches all registration and cleanup methods for packets.
    /// </summary>
    private void CacheRegistrationMethods()
    {
        CacheRegistrationMethods(registerBroadcastPacketListenersMethods, nameof(RegisterBroadcastPacketListenerEvent), broadcastPacketRetrievalSubscriptionDelegates);
        CacheRegistrationMethods(registerEntityPacketListenersMethods, nameof(RegisterEntityPacketListenerEvent), entityPacketRetrievalSubscriptionDelegates);
        CacheRegistrationMethods(clearRoutesMethods, nameof(ClearRouter), uniqueRetrievalSubscriptionDelegates);
    }

    /// <summary>
    /// Creates generic method instances for each packet type listener to be registered into the <see cref="NetworkEntity"/>.
    /// </summary>
    private void CacheRegistrationMethods(
        Dictionary<Type, MethodInfo> listenerRegistrationMethods,
        string methodName,
        List<PacketRetrievalDelegatePair> packetRetrievalDelegates)
    {
        MethodInfo registerPacketMethod = typeof(NetworkManager).GetMethod(methodName, BindingFlags.NonPublic | BindingFlags.Static)!;

        foreach ((Type packetType, Delegate @delegate) in packetRetrievalDelegates)
        {
            MethodInfo genericMethod = registerPacketMethod.MakeGenericMethod(packetType);
            listenerRegistrationMethods.TryAdd(packetType, genericMethod);
        }
    }

    /// <summary>
    /// Caches packet listener methods for all packet listener interfaces.
    /// </summary>
    private void CachePacketListenerMethods()
    {
        CachePacketListenerMethods(clientBroadcastPacketListenerMethods, typeof(IPacketListenerClient<>), nameof(IPacketListenerClient<>.OnClientPacketArrived));
        CachePacketListenerMethods(serverBroadcastPacketListenerMethods, typeof(IPacketListenerServer<>), nameof(IPacketListenerServer<>.OnServerPacketArrived));
        CachePacketListenerMethods(clientEntityPacketListenerMethods, typeof(IPacketListenerClientEntity<>), nameof(IPacketListenerClientEntity<>.OnEntityClientPacketArrived));
        CachePacketListenerMethods(serverEntityPacketListenerMethods, typeof(IPacketListenerServerEntity<>), nameof(IPacketListenerServerEntity<>.OnEntityServerPacketArrived));
    }

    /// <summary>
    /// Discovers all types implementing a given packet listener interface and caches their methods.
    /// </summary>
    private static void CachePacketListenerMethods(
        Dictionary<Type, Dictionary<Type, List<MethodInfo>>> packetListenerMethods,
        Type listenerType,
        string packetListenerMethodName)
    {
        foreach (Type listenerClass in GetGenericsWith(listenerType))
        {
            Dictionary<Type, List<MethodInfo>> listenerMethodDictionary = [];
            packetListenerMethods.Add(listenerClass, listenerMethodDictionary);

            foreach (Type packetListener in GetGenericInterfacesWith(listenerType, listenerClass))
            {
                Type packetType = packetListener.GetGenericArguments().First();

                List<MethodInfo> listenerMethods = [.. packetListener.GetMethods().Where(m => m.Name == packetListenerMethodName)];

                listenerMethodDictionary.Add(packetType, listenerMethods);
            }
        }
    }

    /// <summary>
    /// Finds all types that implement a generic interface.
    /// </summary>
    private static IEnumerable<Type> GetGenericsWith(Type type)
        => AppDomain.CurrentDomain
            .GetAssemblies()
            .SelectMany(a =>
                a.GetTypes().Where(
                    t => t.GetInterfaces().Any(
                        i => i.IsGenericType && i.GetGenericTypeDefinition() == type)));

    // Gets all generic interfaces of a specific definition on a type
    private static IEnumerable<Type> GetGenericInterfacesWith(Type interfaceType, Type type)
        => type.GetInterfaces().Where(i => i.IsGenericType && i.GetGenericTypeDefinition() == interfaceType);

    #endregion

    // Identifies whether packet routing is client-side or server-side
    private enum NetworkType { Client, Server }

    private readonly record struct PacketRetrievalDelegatePair(Type PacketType, Delegate Delegate)
    {
        public static implicit operator (Type packetType, Delegate @delegate)(PacketRetrievalDelegatePair value) => (value.PacketType, value.Delegate);
        public static implicit operator PacketRetrievalDelegatePair((Type packetType, Delegate @delegate) value) => new(value.packetType, value.@delegate);
    }
}