com.alicizax.unity.ui.extension/Runtime/UXComponent/Hotkey/UXHotkeyRegisterManager.cs

#if INPUTSYSTEM_SUPPORT
using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using AlicizaX.UI.Runtime;
using UnityEngine;
using UnityEngine.InputSystem;

namespace UnityEngine.UI
{
    internal enum EHotkeyPressType : byte
    {
        Started = 0,
        Performed = 1
    }

    internal readonly struct HotkeyRegistration
    {
        public readonly IHotkeyTrigger Trigger;
        public readonly EHotkeyPressType PressType;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public HotkeyRegistration(IHotkeyTrigger trigger, EHotkeyPressType pressType)
        {
            Trigger = trigger;
            PressType = pressType;
        }
    }

    internal sealed class HotkeyScope
    {
        public HotkeyScope(UIHolderObjectBase holder)
        {
            Holder = holder;
            HierarchyDepth = GetHierarchyDepth(holder.transform);
            ParentHolder = FindParentHolder(holder);
            BlocksLowerScopes = ParentHolder == null;
        }

        public readonly UIHolderObjectBase Holder;
        public readonly UIHolderObjectBase ParentHolder;
        public readonly int HierarchyDepth;
        public readonly bool BlocksLowerScopes;
        public readonly Dictionary<string, List<HotkeyRegistration>> RegistrationsByAction = new(StringComparer.Ordinal);

        public bool LifecycleActive;
        public ulong ActivationSerial;

        public Action OnBeforeShow;
        public Action OnBeforeClosed;
        public Action OnDestroy;

        private Canvas _canvas;

        public Canvas Canvas
        {
            get
            {
                if (_canvas == null && Holder != null)
                {
                    _canvas = Holder.GetComponent<Canvas>();
                }

                return _canvas;
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static int GetHierarchyDepth(Transform current)
        {
            int depth = 0;
            while (current != null)
            {
                depth++;
                current = current.parent;
            }

            return depth;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static UIHolderObjectBase FindParentHolder(UIHolderObjectBase holder)
        {
            if (holder == null)
            {
                return null;
            }

            Transform current = holder.transform.parent;
            while (current != null)
            {
                if (current.TryGetComponent<UIHolderObjectBase>(out var parentHolder))
                {
                    return parentHolder;
                }

                current = current.parent;
            }

            return null;
        }
    }

    internal sealed class ActionRegistrationBucket
    {
        public InputActionReference ActionReference;
        public Action<InputAction.CallbackContext> StartedHandler;
        public Action<InputAction.CallbackContext> PerformedHandler;
        public int StartedCount;
        public int PerformedCount;

        public int TotalCount => StartedCount + PerformedCount;
    }

    internal readonly struct TriggerRegistration
    {
        public readonly string ActionId;
        public readonly UIHolderObjectBase Holder;
        public readonly EHotkeyPressType PressType;

        public TriggerRegistration(string actionId, UIHolderObjectBase holder, EHotkeyPressType pressType)
        {
            ActionId = actionId;
            Holder = holder;
            PressType = pressType;
        }
    }

    internal static class UXHotkeyRegisterManager
    {
        private static readonly Dictionary<string, ActionRegistrationBucket> _actions = new(StringComparer.Ordinal);
        private static readonly Dictionary<IHotkeyTrigger, TriggerRegistration> _triggerMap = new();
        private static readonly Dictionary<UIHolderObjectBase, HotkeyScope> _scopes = new();
        private static readonly List<HotkeyScope> _leafScopes = new();
        private static readonly HashSet<UIHolderObjectBase> _ancestorHolders = new();

        private static ulong _serialCounter;

#if UNITY_EDITOR
        [UnityEditor.Callbacks.DidReloadScripts]
        internal static void ClearHotkeyRegistry()
        {
            IHotkeyTrigger[] triggers = new IHotkeyTrigger[_triggerMap.Count];
            int index = 0;
            foreach (var kvp in _triggerMap)
            {
                triggers[index++] = kvp.Key;
            }

            for (int i = 0; i < triggers.Length; i++)
            {
                UnregisterHotkey(triggers[i]);
            }

            _actions.Clear();
            _triggerMap.Clear();
            _scopes.Clear();
            _leafScopes.Clear();
            _ancestorHolders.Clear();
            _serialCounter = 0;
        }
#endif

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal static void RegisterHotkey(IHotkeyTrigger trigger, UIHolderObjectBase holder, InputActionReference action, EHotkeyPressType pressType)
        {
            if (trigger == null || holder == null || action == null || action.action == null)
            {
                return;
            }

            UnregisterHotkey(trigger);

            string actionId = action.action.id.ToString();
            HotkeyScope scope = GetOrCreateScope(holder);
            ActionRegistrationBucket bucket = GetOrCreateBucket(actionId, action);
            HotkeyRegistration registration = new HotkeyRegistration(trigger, pressType);

            AdjustBucketSubscription(bucket, pressType, true);
            AddScopeRegistration(scope, actionId, registration);

            if (scope.LifecycleActive)
            {
                scope.ActivationSerial = ++_serialCounter;
            }

            _triggerMap[trigger] = new TriggerRegistration(actionId, holder, pressType);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal static void UnregisterHotkey(IHotkeyTrigger trigger)
        {
            if (trigger == null || !_triggerMap.TryGetValue(trigger, out var triggerRegistration))
            {
                return;
            }

            if (_actions.TryGetValue(triggerRegistration.ActionId, out var bucket))
            {
                RemoveActionRegistration(bucket, triggerRegistration.PressType, triggerRegistration.ActionId);
            }

            if (_scopes.TryGetValue(triggerRegistration.Holder, out var scope))
            {
                RemoveScopeRegistration(scope, triggerRegistration.ActionId, trigger);
                ReleaseScopeIfEmpty(scope);
            }

            _triggerMap.Remove(trigger);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static ActionRegistrationBucket GetOrCreateBucket(string actionId, InputActionReference action)
        {
            if (_actions.TryGetValue(actionId, out var bucket))
            {
                return bucket;
            }

            bucket = new ActionRegistrationBucket
            {
                ActionReference = action,
                StartedHandler = _ => Dispatch(actionId, EHotkeyPressType.Started),
                PerformedHandler = _ => Dispatch(actionId, EHotkeyPressType.Performed)
            };
            _actions[actionId] = bucket;
            return bucket;
        }

        private static HotkeyScope GetOrCreateScope(UIHolderObjectBase holder)
        {
            if (_scopes.TryGetValue(holder, out var scope))
            {
                return scope;
            }

            scope = new HotkeyScope(holder)
            {
                LifecycleActive = IsHolderVisible(holder),
                ActivationSerial = ++_serialCounter
            };

            scope.OnBeforeShow = () => ActivateScope(holder);
            scope.OnBeforeClosed = () => DeactivateScope(holder);
            scope.OnDestroy = () => DestroyScope(holder);

            holder.OnWindowBeforeShowEvent += scope.OnBeforeShow;
            holder.OnWindowBeforeClosedEvent += scope.OnBeforeClosed;
            holder.OnWindowDestroyEvent += scope.OnDestroy;

            _scopes[holder] = scope;
            return scope;
        }

        private static void ActivateScope(UIHolderObjectBase holder)
        {
            if (_scopes.TryGetValue(holder, out var scope))
            {
                scope.LifecycleActive = true;
                scope.ActivationSerial = ++_serialCounter;
            }
        }

        private static void DeactivateScope(UIHolderObjectBase holder)
        {
            if (_scopes.TryGetValue(holder, out var scope))
            {
                scope.LifecycleActive = false;
            }
        }

        private static void DestroyScope(UIHolderObjectBase holder)
        {
            if (holder == null || !_scopes.TryGetValue(holder, out var scope))
            {
                return;
            }

            List<IHotkeyTrigger> triggers = null;
            foreach (var pair in scope.RegistrationsByAction)
            {
                List<HotkeyRegistration> registrations = pair.Value;
                for (int i = 0; i < registrations.Count; i++)
                {
                    triggers ??= new List<IHotkeyTrigger>(registrations.Count);
                    triggers.Add(registrations[i].Trigger);
                }
            }

            if (triggers != null)
            {
                for (int i = 0; i < triggers.Count; i++)
                {
                    UnregisterHotkey(triggers[i]);
                }
            }

            DetachScope(scope);
        }

        private static void DetachScope(HotkeyScope scope)
        {
            if (scope == null || scope.Holder == null)
            {
                return;
            }

            scope.Holder.OnWindowBeforeShowEvent -= scope.OnBeforeShow;
            scope.Holder.OnWindowBeforeClosedEvent -= scope.OnBeforeClosed;
            scope.Holder.OnWindowDestroyEvent -= scope.OnDestroy;
            _scopes.Remove(scope.Holder);
        }

        private static void ReleaseScopeIfEmpty(HotkeyScope scope)
        {
            if (scope != null && scope.RegistrationsByAction.Count == 0)
            {
                DetachScope(scope);
            }
        }

        private static void AddScopeRegistration(HotkeyScope scope, string actionId, HotkeyRegistration registration)
        {
            if (!scope.RegistrationsByAction.TryGetValue(actionId, out var registrations))
            {
                registrations = new List<HotkeyRegistration>();
                scope.RegistrationsByAction[actionId] = registrations;
            }

            registrations.Add(registration);
        }

        private static void RemoveScopeRegistration(HotkeyScope scope, string actionId, IHotkeyTrigger trigger)
        {
            if (!scope.RegistrationsByAction.TryGetValue(actionId, out var registrations))
            {
                return;
            }

            for (int i = registrations.Count - 1; i >= 0; i--)
            {
                if (ReferenceEquals(registrations[i].Trigger, trigger))
                {
                    registrations.RemoveAt(i);
                    break;
                }
            }

            if (registrations.Count == 0)
            {
                scope.RegistrationsByAction.Remove(actionId);
            }
        }

        private static void RemoveActionRegistration(ActionRegistrationBucket bucket, EHotkeyPressType pressType, string actionId)
        {
            AdjustBucketSubscription(bucket, pressType, false);
            if (bucket.TotalCount == 0)
            {
                _actions.Remove(actionId);
            }
        }

        private static void AdjustBucketSubscription(ActionRegistrationBucket bucket, EHotkeyPressType pressType, bool add)
        {
            InputAction inputAction = bucket.ActionReference != null ? bucket.ActionReference.action : null;
            if (inputAction == null)
            {
                return;
            }

            switch (pressType)
            {
                case EHotkeyPressType.Started:
                    if (add)
                    {
                        if (bucket.StartedCount == 0)
                        {
                            inputAction.started += bucket.StartedHandler;
                        }

                        bucket.StartedCount++;
                    }
                    else if (bucket.StartedCount > 0)
                    {
                        bucket.StartedCount--;
                        if (bucket.StartedCount == 0)
                        {
                            inputAction.started -= bucket.StartedHandler;
                        }
                    }

                    break;
                case EHotkeyPressType.Performed:
                    if (add)
                    {
                        if (bucket.PerformedCount == 0)
                        {
                            inputAction.performed += bucket.PerformedHandler;
                        }

                        bucket.PerformedCount++;
                    }
                    else if (bucket.PerformedCount > 0)
                    {
                        bucket.PerformedCount--;
                        if (bucket.PerformedCount == 0)
                        {
                            inputAction.performed -= bucket.PerformedHandler;
                        }
                    }

                    break;
            }

            if (bucket.TotalCount > 0)
            {
                inputAction.Enable();
            }
            else
            {
                inputAction.Disable();
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static void Dispatch(string actionId, EHotkeyPressType pressType)
        {
            HotkeyScope leafScope = GetTopLeafScope();
            if (leafScope == null)
            {
                return;
            }

            TryDispatchToScopeChain(leafScope, actionId, pressType);
        }

        private static bool TryDispatchToScopeChain(HotkeyScope leafScope, string actionId, EHotkeyPressType pressType)
        {
            HotkeyScope current = leafScope;
            while (current != null)
            {
                if (TryGetLatestRegistration(current, actionId, pressType, out var registration))
                {
                    registration.Trigger?.HotkeyActionTrigger();
                    return true;
                }

                UIHolderObjectBase parentHolder = current.ParentHolder;
                current = parentHolder != null && _scopes.TryGetValue(parentHolder, out var parentScope)
                    ? parentScope
                    : null;
            }

            return false;
        }

        private static bool TryGetLatestRegistration(HotkeyScope scope, string actionId, EHotkeyPressType pressType, out HotkeyRegistration registration)
        {
            if (scope.RegistrationsByAction.TryGetValue(actionId, out var registrations))
            {
                for (int i = registrations.Count - 1; i >= 0; i--)
                {
                    HotkeyRegistration candidate = registrations[i];
                    if (candidate.PressType == pressType && candidate.Trigger != null)
                    {
                        registration = candidate;
                        return true;
                    }
                }
            }

            registration = default;
            return false;
        }

        private static HotkeyScope GetTopLeafScope()
        {
            _leafScopes.Clear();
            _ancestorHolders.Clear();

            foreach (var scope in _scopes.Values)
            {
                if (!IsScopeActive(scope)
#if UX_NAVIGATION
                          || !UXNavigationRuntime.IsHolderWithinTopScope(scope.Holder)
#endif
                   )
                {
                    continue;
                }

                UIHolderObjectBase parentHolder = scope.ParentHolder;
                while (parentHolder != null)
                {
                    _ancestorHolders.Add(parentHolder);
                    if (_scopes.TryGetValue(parentHolder, out var parentScope))
                    {
                        parentHolder = parentScope.ParentHolder;
                    }
                    else
                    {
                        break;
                    }
                }
            }

            foreach (var scope in _scopes.Values)
            {
                if (IsScopeActive(scope)
#if UX_NAVIGATION
                    && UXNavigationRuntime.IsHolderWithinTopScope(scope.Holder)
#endif
                    && !_ancestorHolders.Contains(scope.Holder))
                {
                    _leafScopes.Add(scope);
                }
            }

            if (_leafScopes.Count == 0)
            {
                return null;
            }

            _leafScopes.Sort(CompareScopePriority);
            return _leafScopes[0];
        }

        private static bool IsScopeActive(HotkeyScope scope)
        {
            if (scope == null || !scope.LifecycleActive)
            {
                return false;
            }

            UIHolderObjectBase holder = scope.Holder;
            if (holder == null || !holder.IsValid())
            {
                return false;
            }

            if (!holder.gameObject.activeInHierarchy)
            {
                return false;
            }

            Canvas canvas = scope.Canvas;
            return canvas != null && canvas.gameObject.layer == UIComponent.UIShowLayer;
        }

        private static bool IsHolderVisible(UIHolderObjectBase holder)
        {
            if (holder == null || !holder.gameObject.activeInHierarchy)
            {
                return false;
            }

            Canvas canvas = holder.GetComponent<Canvas>();
            return canvas != null && canvas.gameObject.layer == UIComponent.UIShowLayer;
        }

        private static int CompareScopePriority(HotkeyScope left, HotkeyScope right)
        {
            int leftDepth = left.Canvas != null ? left.Canvas.sortingOrder : int.MinValue;
            int rightDepth = right.Canvas != null ? right.Canvas.sortingOrder : int.MinValue;
            int depthCompare = rightDepth.CompareTo(leftDepth);
            if (depthCompare != 0)
            {
                return depthCompare;
            }

            int hierarchyCompare = right.HierarchyDepth.CompareTo(left.HierarchyDepth);
            if (hierarchyCompare != 0)
            {
                return hierarchyCompare;
            }

            return right.ActivationSerial.CompareTo(left.ActivationSerial);
        }

        private static UIHolderObjectBase FindParentHolder(UIHolderObjectBase holder)
        {
            if (holder == null)
            {
                return null;
            }

            Transform current = holder.transform.parent;
            while (current != null)
            {
                if (current.TryGetComponent<UIHolderObjectBase>(out var parentHolder))
                {
                    return parentHolder;
                }

                current = current.parent;
            }

            return null;
        }

#if UNITY_EDITOR
        public static string GetDebugInfo()
        {
            return $"Actions: {_actions.Count}, Triggers: {_triggerMap.Count}, Scopes: {_scopes.Count}";
        }
#endif
    }
}

namespace UnityEngine.UI
{
    public static class UXHotkeyHotkeyExtension
    {
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void BindHotKey(this IHotkeyTrigger trigger)
        {
            if (trigger?.HotkeyAction == null)
            {
                return;
            }

            if (trigger is not Component component)
            {
                return;
            }

            UIHolderObjectBase holder = component.GetComponentInParent<UIHolderObjectBase>(true);
            if (holder == null)
            {
                Debug.LogWarning($"{nameof(HotkeyComponent)} could not find a {nameof(UIHolderObjectBase)} owner.", component);
                return;
            }

            UXHotkeyRegisterManager.RegisterHotkey(trigger, holder, trigger.HotkeyAction, trigger.HotkeyPressType);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void UnBindHotKey(this IHotkeyTrigger trigger)
        {
            if (trigger?.HotkeyAction != null)
            {
                UXHotkeyRegisterManager.UnregisterHotkey(trigger);
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void BindHotKeyBatch(this IHotkeyTrigger[] triggers)
        {
            if (triggers == null)
            {
                return;
            }

            for (int i = 0; i < triggers.Length; i++)
            {
                triggers[i]?.BindHotKey();
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void UnBindHotKeyBatch(this IHotkeyTrigger[] triggers)
        {
            if (triggers == null)
            {
                return;
            }

            for (int i = 0; i < triggers.Length; i++)
            {
                triggers[i]?.UnBindHotKey();
            }
        }
    }
}
#endif