ShrinkNetwork/Runtime/Routing/ShrinkNetworkRegHelper.cs

#nullable enable

using System;
using System.Collections.Generic;
using System.Reflection;
using Cysharp.Threading.Tasks;

namespace ShrinkNetwork
{
    public static class ShrinkNetworkRegHelper
    {
        private static readonly MethodInfo AwaitUniTaskResponseMethod =
            typeof(ShrinkNetworkRegHelper).GetMethod(nameof(AwaitUniTaskResponse), BindingFlags.NonPublic | BindingFlags.Static)!;

        public static void RegisterAttributedMessages(ShrinkNetworkService service)
        {
            RegisterAttributedMessages(service, ShrinkNetworkGeneratedRegistry.GetAttributedMessageTypes());
        }

        public static void RegisterAttributedMessages(ShrinkNetworkService service, IEnumerable<Type> messageTypes)
        {
            if (service == null)
                throw new ArgumentNullException(nameof(service));
            if (messageTypes == null)
                throw new ArgumentNullException(nameof(messageTypes));

            foreach (var type in messageTypes)
            {
                if (type == null || !typeof(IShrinkNetworkMessage).IsAssignableFrom(type))
                    continue;

                var attr = type.GetCustomAttribute<ShrinkNetworkMessageAttribute>(false);
                if (attr == null || service.MessageRegistry.TryGetMeta(type, out _))
                    continue;

                service.RegisterMessage(type, attr.Opcode, attr.Route);
            }
        }

        public static void RegisterStaticHandlers(ShrinkNetworkService service)
            => RegisterStaticHandlers(service, ShrinkNetworkGeneratedRegistry.GetStaticSubscriberTypes());

        public static void RegisterStaticHandlers(ShrinkNetworkService service, IEnumerable<Type> subscriberTypes)
        {
            if (service == null)
                throw new ArgumentNullException(nameof(service));
            if (subscriberTypes == null)
                throw new ArgumentNullException(nameof(subscriberTypes));

            foreach (var type in subscriberTypes)
            {
                if (type == null || type.GetCustomAttribute<ShrinkNetworkSubscriberAttribute>(false) == null)
                    continue;

                ScanMethodsAndRegister(service, null, type,
                    type.GetMethods(BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic));
            }
        }

        public static void RegisterHandlers(ShrinkNetworkService service, object target)
            => RegisterHandlersInternal(service, target);

        private static void RegisterHandlersInternal(ShrinkNetworkService service, object? target)
        {
            if (target == null)
                return;

            var ownerType = target.GetType();
            if (ownerType.GetCustomAttribute<ShrinkNetworkSubscriberAttribute>(false) == null)
                return;

            ScanMethodsAndRegister(service, target, ownerType,
                ownerType.GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic));
        }

        private static void ScanMethodsAndRegister(ShrinkNetworkService service, object? target, Type ownerType, MethodInfo[] methods)
        {
            foreach (var method in methods)
            {
                if (target != null && method.IsStatic)
                    continue;
                var attributes = method.GetCustomAttributes(typeof(ShrinkNetworkSubscribeAttribute), false);
                if (attributes.Length == 0)
                    continue;
                var subscribeAttr = (ShrinkNetworkSubscribeAttribute)attributes[0];
                var requirement = new ShrinkNetworkPermissionRequirement(subscribeAttr.Authority, subscribeAttr.Permission);

                if (!TryParseHandlerSignature(method, out var hasContext, out var messageType, out var responseType))
                {
                    ShrinkNetworkLogger.Warn($"[ShrinkNetwork] Invalid handler signature: {ownerType.FullName}.{method.Name}");
                    continue;
                }

                var resolvedMessageType = messageType!;
                EnsureMessageRegistered(service, resolvedMessageType);
                if (responseType != null)
                    EnsureMessageRegistered(service, responseType);

                if (responseType == null)
                {
                    service.RegisterHandler(resolvedMessageType, (context, message) =>
                        InvokeMessageHandler(target, method, hasContext, context, message), requirement);
                    continue;
                }

                var resolvedResponseType = responseType!;
                service.RegisterRequestHandler(resolvedMessageType, resolvedResponseType, (context, request) =>
                    InvokeRequestHandler(target, method, hasContext, resolvedResponseType, context, request), requirement);
            }
        }

        private static bool TryParseHandlerSignature(MethodInfo method, out bool hasContext, out Type? messageType, out Type? responseType)
        {
            hasContext = false;
            messageType = null;
            responseType = null;

            var parameters = method.GetParameters();
            if (parameters.Length == 1)
            {
                messageType = parameters[0].ParameterType;
            }
            else if (parameters.Length == 2 && parameters[0].ParameterType == typeof(ShrinkNetworkContext))
            {
                hasContext = true;
                messageType = parameters[1].ParameterType;
            }
            else
            {
                return false;
            }

            if (!typeof(IShrinkNetworkMessage).IsAssignableFrom(messageType))
                return false;

            var returnType = method.ReturnType;
            if (returnType == typeof(void) || returnType == typeof(UniTask))
                return true;

            if (typeof(IShrinkNetworkResponse).IsAssignableFrom(returnType))
            {
                responseType = returnType;
                return true;
            }

            if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() == typeof(UniTask<>))
            {
                var resultType = returnType.GetGenericArguments()[0];
                if (typeof(IShrinkNetworkResponse).IsAssignableFrom(resultType))
                {
                    responseType = resultType;
                    return true;
                }
            }

            return false;
        }

        private static void EnsureMessageRegistered(ShrinkNetworkService service, Type messageType)
        {
            if (service.MessageRegistry.TryGetMeta(messageType, out _))
                return;

            var attr = messageType.GetCustomAttribute<ShrinkNetworkMessageAttribute>(false);
            if (attr == null)
                throw new InvalidOperationException($"Message type {messageType.FullName} must declare [ShrinkNetworkMessage].");

            service.RegisterMessage(messageType, attr.Opcode, attr.Route);
        }

        private static async UniTask InvokeMessageHandler(object? target, MethodInfo method, bool hasContext,
            ShrinkNetworkContext context, object message)
        {
            var args = hasContext ? new object?[] { context, message } : new object?[] { message };
            var result = method.Invoke(target, args);
            if (method.ReturnType == typeof(UniTask))
                await (UniTask)(result ?? throw new InvalidOperationException($"Handler returned null UniTask: {method.DeclaringType?.FullName}.{method.Name}"));
        }

        private static async UniTask<object?> InvokeRequestHandler(object? target, MethodInfo method, bool hasContext,
            Type responseType, ShrinkNetworkContext context, object request)
        {
            var args = hasContext ? new object?[] { context, request } : new object?[] { request };
            var result = method.Invoke(target, args);
            if (result == null)
                return null;

            if (responseType.IsInstanceOfType(result))
                return result;

            if (method.ReturnType.IsGenericType && method.ReturnType.GetGenericTypeDefinition() == typeof(UniTask<>))
                return await AwaitUniTaskResponseObject(result, responseType);

            return result;
        }

        private static UniTask<object?> AwaitUniTaskResponseObject(object taskObject, Type responseType)
        {
            return (UniTask<object?>)AwaitUniTaskResponseMethod
                .MakeGenericMethod(responseType)
                .Invoke(null, new[] { taskObject })!;
        }

        private static async UniTask<object?> AwaitUniTaskResponse<TResponse>(UniTask<TResponse> task)
            where TResponse : class, IShrinkNetworkResponse
        {
            return await task;
        }
    }
}