ShrinkNetwork.Integration.EventBus/ShrinkNetworkEventBusBridge.cs

#nullable enable
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using Cysharp.Threading.Tasks;
using ShrinkEventBus;
using UnityEngine;

namespace ShrinkNetwork.Integration
{
    public static class ShrinkNetworkEventBusBridge
    {
        private sealed class ServiceRegistration
        {
            public ShrinkNetworkService Service = null!;
            public Func<ShrinkNetworkSession, IShrinkEvent, bool>? SessionFilter;
            public Dictionary<string, long> AppliedDeltaVersions { get; } = new();
        }

        private static readonly object SyncRoot = new();
        private static readonly Dictionary<ShrinkNetworkService, ServiceRegistration> RegisteredServices = new();
        private static readonly AsyncLocal<int> SuppressForwardDepth = new();
        private static Dictionary<Type, ShrinkNetworkEventRegistration> _networkEventTypes = new();
        private static bool _initialized;

        [RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.SubsystemRegistration)]
        private static void ResetStaticState()
        {
            EventBus.Posted -= HandleLocalEventPosted;
            lock (SyncRoot)
            {
                RegisteredServices.Clear();
                FallbackDeltaVersions.Clear();
            }

            _networkEventTypes = new Dictionary<Type, ShrinkNetworkEventRegistration>();
            _initialized = false;
        }

        private static void EnsureBridgeInitialized()
        {
            if (_initialized)
                return;

            _initialized = true;
            RefreshNetworkEventTypes();
            EventBus.Posted += HandleLocalEventPosted;
        }

        public static void RegisterService(ShrinkNetworkService service,
            Func<ShrinkNetworkSession, IShrinkEvent, bool>? sessionFilter = null)
        {
            if (service == null)
                throw new ArgumentNullException(nameof(service));

            EnsureBridgeInitialized();
            RefreshNetworkEventTypes();

            lock (SyncRoot)
            {
                if (RegisteredServices.TryGetValue(service, out var existing))
                {
                    existing.SessionFilter = sessionFilter;
                    return;
                }

                RegisteredServices.Add(service, new ServiceRegistration
                {
                    Service = service,
                    SessionFilter = sessionFilter
                });
            }

            RegisterInboundBridgeHandlers(service);
        }

        public static void UnregisterService(ShrinkNetworkService service)
        {
            if (service == null)
                return;

            lock (SyncRoot)
            {
                RegisteredServices.Remove(service);
            }
        }

        public static void RefreshBindings()
        {
            RefreshNetworkEventTypes();

            List<ShrinkNetworkService> services;
            lock (SyncRoot)
            {
                services = RegisteredServices.Keys.ToList();
            }

            foreach (var service in services)
            {
                RegisterInboundBridgeHandlers(service);
            }
        }

        public static UniTask PublishAsync<TEvent>(ShrinkNetworkSession session, TEvent eventArgs, string? route = null)
            where TEvent : IShrinkEvent, IShrinkNetworkMessage
        {
            if (session == null)
                throw new ArgumentNullException(nameof(session));
            if (eventArgs == null)
                throw new ArgumentNullException(nameof(eventArgs));

            EnsureNetworkEventType(typeof(TEvent));
            EnsurePublishableEventType(typeof(TEvent));
            return session.SendAsync(eventArgs, route);
        }

        public static async UniTask PublishAsync<TEvent>(IEnumerable<ShrinkNetworkSession> sessions, TEvent eventArgs,
            string? route = null)
            where TEvent : IShrinkEvent, IShrinkNetworkMessage
        {
            if (sessions == null)
                throw new ArgumentNullException(nameof(sessions));
            if (eventArgs == null)
                throw new ArgumentNullException(nameof(eventArgs));

            EnsureNetworkEventType(typeof(TEvent));
            EnsurePublishableEventType(typeof(TEvent));

            var groupedSessions = sessions
                .Where(session => session != null)
                .GroupBy(session => session.Service);

            foreach (var serviceGroup in groupedSessions)
            {
                var sessionList = serviceGroup.ToList();
                if (sessionList.Count == 0)
                    continue;

                var payload = serviceGroup.Key.Serializer.Serialize(eventArgs);

                foreach (var session in sessionList)
                    await serviceGroup.Key.SendSerializedAsync(session, typeof(TEvent), payload, route);
            }
        }

        public static async UniTask<ShrinkNetworkEventResultResponse> RequestResultAsync<TEvent>(ShrinkNetworkSession session,
            TEvent eventArgs, ShrinkRpcCallOptions? options = null)
            where TEvent : IShrinkResultEvent<EventResult>, IShrinkNetworkRequest
        {
            if (session == null)
                throw new ArgumentNullException(nameof(session));
            if (eventArgs == null)
                throw new ArgumentNullException(nameof(eventArgs));

            EnsureResultRequestEventType(typeof(TEvent));

            options ??= new ShrinkRpcCallOptions();
            if (string.IsNullOrWhiteSpace(options.DebugLabel))
                options.DebugLabel = typeof(TEvent).Name;

            var response = await session.RpcAsync<TEvent, ShrinkNetworkEventResultResponse>(eventArgs,
                options);

            if (response.IsSuccess)
            {
                eventArgs.SetResult(response.Result);
                if (eventArgs is IShrinkCancelableEvent cancelable)
                    cancelable.SetCanceled(response.IsCanceled);
            }

            return response;
        }

        private static void HandleLocalEventPosted(IShrinkEvent eventArgs, Type eventType, ShrinkBusKey busKey)
        {
            if (eventArgs == null || eventType == null)
                return;
            if (SuppressForwardDepth.Value > 0)
                return;
            if (!TryGetNetworkEventMeta(eventType, out _))
                return;
            if (eventArgs is IShrinkNetworkRequest)
                return;
            if (eventArgs is not IShrinkNetworkMessage networkMessage)
                return;

            ForwardEventAsync(eventArgs, eventType, networkMessage).Forget();
        }

        private static async UniTaskVoid ForwardEventAsync(IShrinkEvent eventArgs, Type eventType,
            IShrinkNetworkMessage networkMessage)
        {
            List<ServiceRegistration> registrations;
            lock (SyncRoot)
            {
                registrations = RegisteredServices.Values.ToList();
            }

            var batches = new List<(ServiceRegistration Registration, List<ShrinkNetworkSession> Sessions, byte[] Payload)>();
            foreach (var registration in registrations)
            {
                var sessions = registration.Service.Sessions.Values
                    .Where(session => registration.SessionFilter == null || registration.SessionFilter(session, eventArgs))
                    .ToList();
                if (sessions.Count == 0)
                    continue;

                byte[] payload;
                try
                {
                    // Event payload is immutable after serialization and can be
                    // reused for every session owned by this service.
                    payload = registration.Service.Serializer.Serialize(networkMessage);
                }
                catch (Exception ex)
                {
                    Debug.LogException(ex);
                    Debug.LogError(
                        $"[ShrinkNetwork.Integration] 序列化网络事件失败: {eventArgs.GetType().FullName}");
                    continue;
                }

                // Complete every serializer pass before the first await so mutable
                // event payloads cannot change between target sessions.
                batches.Add((registration, sessions, payload));
            }

            foreach (var batch in batches)
            {
                foreach (var session in batch.Sessions)
                {
                    try
                    {
                        await batch.Registration.Service.SendSerializedAsync(session, eventType, batch.Payload);
                    }
                    catch (Exception ex)
                    {
                        Debug.LogException(ex);
                        Debug.LogError(
                            $"[ShrinkNetwork.Integration] 转发网络事件失败: {eventArgs.GetType().FullName}, Session={session.SessionId}");
                    }
                }
            }
        }

        private static void RegisterInboundBridgeHandlers(ShrinkNetworkService service)
        {
            foreach (var pair in _networkEventTypes)
            {
                if (!service.MessageRegistry.TryGetMeta(pair.Key, out _))
                {
                    service.RegisterMessage(pair.Key, pair.Value.Opcode, pair.Value.Route);
                }

                if (!service.MessageRegistry.TryGetMeta(typeof(ShrinkNetworkEventResultResponse), out _))
                {
                    service.RegisterMessage(typeof(ShrinkNetworkEventResultResponse), -300001,
                        "__integration/event_result_response");
                }

                var isRequestEvent = typeof(IShrinkNetworkRequest).IsAssignableFrom(pair.Key);

                try
                {
                    if (isRequestEvent)
                    {
                        service.RegisterRequestHandler(pair.Key, typeof(ShrinkNetworkEventResultResponse),
                            DispatchIncomingRequestAsync);
                    }
                    else
                    {
                        service.RegisterHandler(pair.Key, DispatchIncomingAsync);
                    }
                }
                catch (InvalidOperationException)
                {
                    // 同一个 service 上该事件类型只能绑定一个处理器。已存在时认为调用方自行接管。
                }
            }
        }

        private static UniTask DispatchIncomingAsync(ShrinkNetworkContext context, object message)
        {
            if (message is not IShrinkEvent eventArgs)
                return UniTask.CompletedTask;

            return _networkEventTypes.TryGetValue(eventArgs.GetType(), out var registration)
                ? registration.Dispatch(eventArgs)
                : UniTask.CompletedTask;
        }

        private static UniTask<object?> DispatchIncomingRequestAsync(ShrinkNetworkContext context, object message)
        {
            if (message is not IShrinkEvent eventArgs)
            {
                return UniTask.FromResult<object?>(new ShrinkNetworkEventResultResponse
                {
                    ErrorCode = ShrinkRpcErrorCode.InvalidResponse,
                    ErrorMessage = "Incoming network event request does not implement IShrinkEvent."
                });
            }

            return _networkEventTypes.TryGetValue(eventArgs.GetType(), out var registration) &&
                   registration.RequestDispatch != null
                ? registration.RequestDispatch(context, eventArgs)
                : UniTask.FromResult<object?>(new ShrinkNetworkEventResultResponse
                {
                    ErrorCode = ShrinkRpcErrorCode.InvalidResponse,
                    ErrorMessage = $"No generated network event request binding for {eventArgs.GetType().FullName}."
                });
        }

        internal static async UniTask DispatchGeneratedAsync<TEvent>(TEvent eventArgs)
            where TEvent : IShrinkEvent
        {
            SuppressForwardDepth.Value++;
            try
            {
                if (TryGetDeltaEvent(eventArgs, typeof(TEvent), out var deltaEvent) &&
                    !TryMarkIncomingDelta(typeof(TEvent), deltaEvent, null))
                {
                    return;
                }

                await EventBus.PostAsync(eventArgs);
            }
            finally
            {
                SuppressForwardDepth.Value--;
            }
        }

        internal static async UniTask<object?> DispatchGeneratedRequestAsync<TEvent>(
            ShrinkNetworkContext context, TEvent eventArgs)
            where TEvent : IShrinkEvent
        {
            SuppressForwardDepth.Value++;
            try
            {
                if (TryGetDeltaEvent(eventArgs, typeof(TEvent), out var deltaEvent) &&
                    !TryMarkIncomingDelta(typeof(TEvent), deltaEvent, context))
                {
                    return BuildResultResponse(eventArgs);
                }

                await EventBus.PostAsync(eventArgs);
                return BuildResultResponse(eventArgs);
            }
            catch (Exception ex)
            {
                Debug.LogException(ex);
                return new ShrinkNetworkEventResultResponse
                {
                    ErrorCode = ShrinkRpcErrorCode.HandlerException,
                    ErrorMessage = ex.Message,
                    Result = eventArgs is IShrinkResultEvent<EventResult> resultEvent
                        ? resultEvent.Result
                        : EventResult.DEFAULT,
                    IsCanceled = eventArgs is IShrinkCancelableEvent cancelable && cancelable.IsCanceled
                };
            }
            finally
            {
                SuppressForwardDepth.Value--;
            }
        }

        private static void RefreshNetworkEventTypes()
        {
            var eventTypes = new Dictionary<Type, ShrinkNetworkEventRegistration>();
            foreach (var registration in ShrinkNetworkEventRegistry.Snapshot())
                eventTypes[registration.EventType] = registration;

            _networkEventTypes = eventTypes;
        }

        private static bool TryGetNetworkEventMeta(Type eventType,
            out ShrinkNetworkEventRegistration registration)
        {
            if (eventType == null)
            {
                registration = default;
                return false;
            }

            return _networkEventTypes.TryGetValue(eventType, out registration);
        }

        private static void EnsureNetworkEventType(Type eventType)
        {
            if (!TryGetNetworkEventMeta(eventType, out _))
                throw new InvalidOperationException(
                    $"Event type {eventType.FullName} must implement IShrinkEvent + IShrinkNetworkMessage and declare [ShrinkNetworkEvent] + [ShrinkNetworkMessage].");
        }

        private static void EnsurePublishableEventType(Type eventType)
        {
            if (typeof(IShrinkNetworkRequest).IsAssignableFrom(eventType))
                throw new InvalidOperationException(
                    $"Event type {eventType.FullName} is a request event. Use RequestResultAsync instead of PublishAsync.");
        }

        private static void EnsureResultRequestEventType(Type eventType)
        {
            EnsureNetworkEventType(eventType);

            if (!typeof(IShrinkNetworkRequest).IsAssignableFrom(eventType))
                throw new InvalidOperationException(
                    $"Event type {eventType.FullName} must implement IShrinkNetworkRequest to use RequestResultAsync.");
            if (!typeof(IShrinkResultEvent<EventResult>).IsAssignableFrom(eventType))
                throw new InvalidOperationException(
                    $"Event type {eventType.FullName} must implement IShrinkResultEvent<EventResult> to use RequestResultAsync.");
        }

        private static ShrinkNetworkEventResultResponse BuildResultResponse(IShrinkEvent eventArgs)
        {
            return new ShrinkNetworkEventResultResponse
            {
                Result = eventArgs is IShrinkResultEvent<EventResult> resultEvent
                    ? resultEvent.Result
                    : EventResult.DEFAULT,
                IsCanceled = eventArgs is IShrinkCancelableEvent cancelable && cancelable.IsCanceled
            };
        }

        private static bool TryGetDeltaEvent(IShrinkEvent eventArgs, Type eventType, out IShrinkNetworkDeltaEvent deltaEvent)
        {
            if (eventArgs is IShrinkNetworkDeltaEvent eventDelta)
            {
                deltaEvent = eventDelta;
                return true;
            }

            deltaEvent = null!;
            return false;
        }

        private static bool TryMarkIncomingDelta(Type eventType, IShrinkNetworkDeltaEvent deltaEvent,
            ShrinkNetworkContext? context)
        {
            if (deltaEvent == null)
                return true;

            var scopeKey = BuildDeltaScopeKey(eventType, deltaEvent, context);

            lock (SyncRoot)
            {
                if (context != null && RegisteredServices.TryGetValue(context.Service, out var serviceRegistration))
                {
                    return TryMarkIncomingDeltaCore(serviceRegistration.AppliedDeltaVersions, scopeKey,
                        deltaEvent.DeltaVersion);
                }

                return TryMarkIncomingDeltaCore(FallbackDeltaVersions, scopeKey, deltaEvent.DeltaVersion);
            }
        }

        private static bool TryMarkIncomingDeltaCore(Dictionary<string, long> versionMap, string scopeKey, long deltaVersion)
        {
            if (versionMap.TryGetValue(scopeKey, out var existingVersion) && deltaVersion <= existingVersion)
                return false;

            versionMap[scopeKey] = deltaVersion;
            return true;
        }

        private static string BuildDeltaScopeKey(Type eventType, IShrinkNetworkDeltaEvent deltaEvent,
            ShrinkNetworkContext? context)
        {
            var sessionId = context?.Session?.SessionId ?? 0;
            var key = string.IsNullOrWhiteSpace(deltaEvent.DeltaKey) ? "__default" : deltaEvent.DeltaKey.Trim();
            return $"{sessionId}:{eventType.FullName}:{key}";
        }

        private static readonly Dictionary<string, long> FallbackDeltaVersions = new();
    }

    /// <summary>
    /// Switches network callbacks to Unity's main thread before entering EventBus
    /// handlers. Use ShrinkNetworkDispatchQueue when a hard per-frame budget is
    /// required instead of an unbounded PlayerLoop backlog.
    /// </summary>
    public sealed class ShrinkNetworkUnityMainThreadDispatchScheduler : IShrinkNetworkDispatchScheduler
    {
        public async UniTask<bool> ScheduleAsync(Func<UniTask> callback)
        {
            if (callback == null)
                throw new ArgumentNullException(nameof(callback));

            await UniTask.SwitchToMainThread();
            await callback();
            return true;
        }
    }
}