#if INPUTSYSTEM_SUPPORT
using System;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.InputSystem;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine.UI;

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

    [StructLayout(LayoutKind.Sequential, Pack = 4)]
    internal struct HotkeyRegistration
    {
        public IHotkeyTrigger button;
        public EHotkeyPressType pressType;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public HotkeyRegistration(IHotkeyTrigger btn, EHotkeyPressType pressType)
        {
            this.button = btn;
            this.pressType = pressType;
        }
    }

    [StructLayout(LayoutKind.Sequential)]
    internal struct HandlerInfo
    {
        public Action<InputAction.CallbackContext> handler;
        public InputActionReference action;

        public HandlerInfo(Action<InputAction.CallbackContext> handler, InputActionReference action)
        {
            this.handler = handler;
            this.action = action;
        }
    }

    internal static class UXHotkeyRegisterManager
    {
        private const int INITIAL_CAPACITY = 32;
        private const int MAX_REGISTRATIONS_PER_ACTION = 16;

        private static HotkeyRegistration[][] _registrationPool;
        private static int[] _registrationCounts;
        private static string[] _actionIds;
        private static int _actionCount;
        private static int _actionCapacity;

        private static HandlerInfo[] _handlers;
        private static InputActionReference[] _actionRefs;

        private static IHotkeyTrigger[] _buttons;
        private static int[] _buttonToActionIndex;
        private static int _buttonCount;
        private static int _buttonCapacity;

        private static Action<InputAction.CallbackContext>[] _cachedHandlers;
        private static int _cachedHandlerCount;

        static UXHotkeyRegisterManager()
        {
            _actionCapacity = INITIAL_CAPACITY;
            _registrationPool = new HotkeyRegistration[_actionCapacity][];
            _registrationCounts = new int[_actionCapacity];
            _actionIds = new string[_actionCapacity];
            _handlers = new HandlerInfo[_actionCapacity];
            _actionRefs = new InputActionReference[_actionCapacity];

            _buttonCapacity = 64;
            _buttons = new IHotkeyTrigger[_buttonCapacity];
            _buttonToActionIndex = new int[_buttonCapacity];

            _cachedHandlers = new Action<InputAction.CallbackContext>[INITIAL_CAPACITY];

            for (int i = 0; i < _actionCapacity; i++)
            {
                _registrationPool[i] = new HotkeyRegistration[MAX_REGISTRATIONS_PER_ACTION];
                _buttonToActionIndex[i] = -1;
            }
        }

#if UNITY_EDITOR
        [UnityEditor.Callbacks.DidReloadScripts]
        internal static void ClearHotkeyRegistry()
        {
            for (int i = 0; i < _buttonCount; i++)
            {
                if (_buttons[i] != null)
                {
                    UnregisterHotkey(_buttons[i]);
                }
            }

            Array.Clear(_registrationCounts, 0, _actionCount);
            Array.Clear(_actionIds, 0, _actionCount);
            Array.Clear(_handlers, 0, _actionCount);
            Array.Clear(_actionRefs, 0, _actionCount);
            Array.Clear(_buttons, 0, _buttonCount);
            Array.Clear(_buttonToActionIndex, 0, _buttonCount);

            _actionCount = 0;
            _buttonCount = 0;
            _cachedHandlerCount = 0;
        }
#endif

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

            string actionId = action.action.id.ToString();
            int actionIndex = FindOrCreateActionIndex(actionId);

            ref int count = ref _registrationCounts[actionIndex];
            if (count >= MAX_REGISTRATIONS_PER_ACTION)
            {
                Debug.LogWarning($"Max registrations reached for action: {actionId}");
                return;
            }

            _registrationPool[actionIndex][count] = new HotkeyRegistration(button, pressType);
            count++;

            int buttonIndex = FindOrCreateButtonIndex(button);
            _buttonToActionIndex[buttonIndex] = actionIndex;

            if (count == 1)
            {
                Action<InputAction.CallbackContext> handler = GetOrCreateHandler(actionIndex);
                _handlers[actionIndex] = new HandlerInfo(handler, action);
                _actionRefs[actionIndex] = action;

                switch (pressType)
                {
                    case EHotkeyPressType.Started:
                        action.action.started += handler;
                        break;
                    case EHotkeyPressType.Performed:
                        action.action.performed += handler;
                        break;
                }

                action.action.Enable();
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void UnregisterHotkey(IHotkeyTrigger button)
        {
            if (button == null)
                return;

            int buttonIndex = FindButtonIndex(button);
            if (buttonIndex < 0)
                return;

            int actionIndex = _buttonToActionIndex[buttonIndex];
            if (actionIndex < 0)
                return;

            ref int count = ref _registrationCounts[actionIndex];
            HotkeyRegistration[] registrations = _registrationPool[actionIndex];

            for (int i = count - 1; i >= 0; i--)
            {
                if (ReferenceEquals(registrations[i].button, button))
                {
                    EHotkeyPressType pressType = registrations[i].pressType;

                    registrations[i] = registrations[count - 1];
                    registrations[count - 1] = default;
                    count--;

                    if (count == 0)
                    {
                        ref HandlerInfo handlerInfo = ref _handlers[actionIndex];
                        InputActionReference actionRef = _actionRefs[actionIndex];

                        if (actionRef != null && actionRef.action != null)
                        {
                            switch (pressType)
                            {
                                case EHotkeyPressType.Started:
                                    actionRef.action.started -= handlerInfo.handler;
                                    break;
                                case EHotkeyPressType.Performed:
                                    actionRef.action.performed -= handlerInfo.handler;
                                    break;
                            }

                            actionRef.action.Disable();
                        }

                        handlerInfo = default;
                        _actionRefs[actionIndex] = null;
                        _actionIds[actionIndex] = null;
                    }

                    break;
                }
            }

            _buttons[buttonIndex] = null;
            _buttonToActionIndex[buttonIndex] = -1;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static Action<InputAction.CallbackContext> GetOrCreateHandler(int actionIndex)
        {
            if (actionIndex < _cachedHandlers.Length && _cachedHandlers[actionIndex] != null)
            {
                return _cachedHandlers[actionIndex];
            }

            if (actionIndex >= _cachedHandlers.Length)
            {
                Array.Resize(ref _cachedHandlers, actionIndex + 1);
            }

            int capturedIndex = actionIndex;
            _cachedHandlers[actionIndex] = ctx => OnHotkeyTriggered(capturedIndex);
            return _cachedHandlers[actionIndex];
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static void OnHotkeyTriggered(int actionIndex)
        {
            int count = _registrationCounts[actionIndex];
            if (count > 0)
            {
                ref HotkeyRegistration registration = ref _registrationPool[actionIndex][count - 1];
                registration.button?.HotkeyActionTrigger();
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static int FindOrCreateActionIndex(string actionId)
        {
            for (int i = 0; i < _actionCount; i++)
            {
                if (_actionIds[i] == actionId)
                    return i;
            }

            if (_actionCount >= _actionCapacity)
            {
                ExpandActionCapacity();
            }

            int index = _actionCount++;
            _actionIds[index] = actionId;
            return index;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static int FindOrCreateButtonIndex(IHotkeyTrigger button)
        {
            for (int i = 0; i < _buttonCount; i++)
            {
                if (ReferenceEquals(_buttons[i], button))
                    return i;
            }

            for (int i = 0; i < _buttonCapacity; i++)
            {
                if (_buttons[i] == null)
                {
                    _buttons[i] = button;
                    if (i >= _buttonCount)
                        _buttonCount = i + 1;
                    return i;
                }
            }

            if (_buttonCount >= _buttonCapacity)
            {
                ExpandButtonCapacity();
            }

            int index = _buttonCount++;
            _buttons[index] = button;
            return index;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static int FindButtonIndex(IHotkeyTrigger button)
        {
            for (int i = 0; i < _buttonCount; i++)
            {
                if (ReferenceEquals(_buttons[i], button))
                    return i;
            }

            return -1;
        }

        private static void ExpandActionCapacity()
        {
            int newCapacity = _actionCapacity * 2;
            Array.Resize(ref _registrationPool, newCapacity);
            Array.Resize(ref _registrationCounts, newCapacity);
            Array.Resize(ref _actionIds, newCapacity);
            Array.Resize(ref _handlers, newCapacity);
            Array.Resize(ref _actionRefs, newCapacity);

            for (int i = _actionCapacity; i < newCapacity; i++)
            {
                _registrationPool[i] = new HotkeyRegistration[MAX_REGISTRATIONS_PER_ACTION];
            }

            _actionCapacity = newCapacity;
        }

        private static void ExpandButtonCapacity()
        {
            int newCapacity = _buttonCapacity * 2;
            Array.Resize(ref _buttons, newCapacity);
            Array.Resize(ref _buttonToActionIndex, newCapacity);

            for (int i = _buttonCapacity; i < newCapacity; i++)
            {
                _buttonToActionIndex[i] = -1;
            }

            _buttonCapacity = newCapacity;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void RegisterHotkeyBatch(Span<IHotkeyTrigger> buttons, InputActionReference action, EHotkeyPressType pressType)
        {
            for (int i = 0; i < buttons.Length; i++)
            {
                RegisterHotkey(buttons[i], action, pressType);
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static void UnregisterHotkeyBatch(Span<IHotkeyTrigger> buttons)
        {
            for (int i = 0; i < buttons.Length; i++)
            {
                UnregisterHotkey(buttons[i]);
            }
        }

#if UNITY_EDITOR
        public static string GetDebugInfo()
        {
            return $"Actions: {_actionCount}/{_actionCapacity}, Buttons: {_buttonCount}/{_buttonCapacity}";
        }
#endif
    }
}

public static class UXHotkeyHotkeyExtension
{
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static void BindHotKey(this IHotkeyTrigger button)
    {
        if (button?.HotkeyAction != null)
        {
            UXHotkeyRegisterManager.RegisterHotkey(
                button,
                button.HotkeyAction,
                button.HotkeyPressType
            );
        }
    }

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

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

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

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

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