continuity/Rai.PacketMediator/PacketDistributor.cs

// Copyright (c) 2023 Timothy Schenk. Subject to the GNU AGPL Version 3 License.

using System.Collections.Concurrent;
using System.Collections.Immutable;
using System.Reflection;
using System.Threading.Channels;
using DotNext.Collections.Generic;
using DotNext.Linq.Expressions;
using DotNext.Metaprogramming;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.VisualBasic.CompilerServices;

namespace Rai.PacketMediator;

public class PacketDistributor<TPacketIdEnum, TSession> where TPacketIdEnum : Enum
{
    private readonly Channel<ValueTuple<byte[], TPacketIdEnum, TSession>> _channel;

    private readonly Assembly[] _sourcesContainingPackets;
    private readonly ConcurrentDictionary<TPacketIdEnum, IPacketHandler<TSession>?> _packetHandlersInstantiation;

    private readonly ImmutableDictionary<TPacketIdEnum,
        Func<byte[], IIncomingPacket>> _deserializationMap;

    public PacketDistributor(IServiceProvider serviceProvider,
        Assembly[] sourcesContainingPackets)
    {
        _channel = Channel.CreateUnbounded<ValueTuple<byte[], TPacketIdEnum, TSession>>(new UnboundedChannelOptions
        {
            AllowSynchronousContinuations = false,
            SingleReader = false,
            SingleWriter = false
        });
        _sourcesContainingPackets = sourcesContainingPackets;
        var packetDictionary = GetAllPackets();

        if (packetDictionary is { Count: 0 })
        {
            throw new IncompleteInitialization();
        }

        var packetHandlers = GetAllPacketHandlersWithId();

        if (packetHandlers is { Count: 0 })
        {
            throw new IncompleteInitialization();
        }

        var tempDeserializationMap =
            new Dictionary<TPacketIdEnum, Func<byte[], IIncomingPacket>>();
        _packetHandlersInstantiation = new ConcurrentDictionary<TPacketIdEnum, IPacketHandler<TSession>?>();
        packetHandlers.ForEach(x =>
        {
            var packetHandler =
                ActivatorUtilities.GetServiceOrCreateInstance(serviceProvider,
                    x.Value);
            _packetHandlersInstantiation.TryAdd(x.Key, packetHandler as IPacketHandler<TSession>);
        });
        foreach (var packetsType in packetDictionary)
        {
            var lambda = CodeGenerator.Lambda<Func<byte[], IIncomingPacket>>(fun =>
            {
                var argPacketData = fun[0];
                var newPacket = packetsType.Value.New();

                var packetVariable = CodeGenerator.DeclareVariable(packetsType.Value, "packet");
                CodeGenerator.Assign(packetVariable, newPacket);
                CodeGenerator.Call(packetVariable, nameof(IIncomingPacket.Deserialize), argPacketData);

                CodeGenerator.Return(packetVariable);
            }).Compile();
            tempDeserializationMap.Add(packetsType.Key, lambda);
        }

        _deserializationMap = tempDeserializationMap.ToImmutableDictionary();
    }

    private Dictionary<TPacketIdEnum, Type> GetAllPackets()
    {
        // ! : item.Attribute cannot be null due to previous Where check
        var packetsWithId = this._sourcesContainingPackets.SelectMany(a => a.GetTypes()
            .Where(type => type is { IsInterface: false, IsAbstract: false } &&
                           type.GetInterfaces().Contains(typeof(IIncomingPacket)))
            .Select(type =>
                new { Type = type, Attribute = type.GetCustomAttribute<PacketIdAttribute<TPacketIdEnum>>() })
            .Where(item => item.Attribute is not null)
            .ToDictionary(item => item.Attribute!.Code, item => item.Type)).ToDictionary();

        return packetsWithId;
    }

    private Dictionary<TPacketIdEnum, Type> GetAllPacketHandlersWithId()
    {
        var packetHandlersWithId = this._sourcesContainingPackets.SelectMany(assembly => assembly.GetTypes()
            .Where(t =>
                t is { IsClass: true, IsAbstract: false } && Array.Exists(t
                    .GetInterfaces(), i =>
                    i.IsGenericType && i.GetGenericTypeDefinition() == typeof(IPacketHandler<TSession>)))
            .Select(type => new
            {
                Type = type,
                PacketId = type
                    .GetInterfaces().First(t1 =>
                        t1 is { IsGenericType: true } &&
                        t1.GetGenericTypeDefinition() == typeof(IPacketHandler<TSession>))
                    .GetGenericArguments().First(t => t.GetCustomAttributes<PacketIdAttribute<TPacketIdEnum>>().Any())
                    .GetCustomAttributes<PacketIdAttribute<TPacketIdEnum>>().First().Code
            })
            .ToDictionary(
                x => x.PacketId, x => x.Type
            )).ToDictionary();

        return packetHandlersWithId;
    }

    public async Task AddPacketAsync(byte[] packetData, TPacketIdEnum operationCode, TSession session)
    {
        await _channel.Writer.WriteAsync((packetData, operationCode, session));
    }

    public async Task DequeueRawPacketAsync(CancellationToken cancellationToken)
    {
        while (await _channel.Reader.WaitToReadAsync(cancellationToken))
        {
            while (_channel.Reader.TryRead(out var item))
            {
                await InvokePacketHandlerAsync(item, cancellationToken);
            }
        }
    }

    private async Task InvokePacketHandlerAsync((byte[], TPacketIdEnum, TSession) valueTuple,
        CancellationToken cancellationToken)
    {
        var (packetData, operationCode, session) = valueTuple;
        if (!_deserializationMap.TryGetValue(operationCode, out var func))
        {
            return;
        }

        var packet = func(packetData);

        // ! I don't see how it's possibly null here.
        await _packetHandlersInstantiation[operationCode]?.TryHandleAsync(packet, session, cancellationToken)!;
    }
}