#if INPUTSYSTEM_SUPPORT
using UnityEngine;
using UnityEngine.InputSystem;
using System.Collections.Generic;
using System.Collections;

internal enum EHotkeyPressType
{
    Started,
    Performed
}

internal static class UXHotkeyComponent
{
    private readonly struct HotkeyRegistration
    {
        public readonly EHotkeyPressType pressType;
        public readonly UXButton button;

        public HotkeyRegistration(UXButton btn, EHotkeyPressType pressType)
        {
            button = btn;
            this.pressType = pressType;
        }
    }

    private static readonly Dictionary<string, List<HotkeyRegistration>> _hotkeyRegistry =
        new Dictionary<string, List<HotkeyRegistration>>(32);


    private static readonly Dictionary<string, (System.Action<InputAction.CallbackContext> handler, InputActionReference action)> _sharedHandlers =
        new Dictionary<string, (System.Action<InputAction.CallbackContext>, InputActionReference)>(32);


    private static readonly Dictionary<UXButton, string> _buttonRegistrations =
        new Dictionary<UXButton, string>(64);


#if UNITY_EDITOR
    [UnityEditor.Callbacks.DidReloadScripts]
    internal static void ClearHotkeyRegistry()
    {
        foreach (var key in _buttonRegistrations.Keys)
        {
            UnregisterHotkey(key);
        }

        _sharedHandlers.Clear();
        _hotkeyRegistry.Clear();
        _buttonRegistrations.Clear();
    }
#endif

    internal static void RegisterHotkey(UXButton button, InputActionReference action, EHotkeyPressType pressType)
    {
        if (action == null || action.action == null || button == null)
            return;

        string actionId = action.action.id.ToString();


        HotkeyRegistration registration = new HotkeyRegistration(button, pressType);


        if (!_hotkeyRegistry.TryGetValue(actionId, out var registrations))
        {
            registrations = new List<HotkeyRegistration>(4);
            _hotkeyRegistry[actionId] = registrations;
        }

        registrations.Add(registration);


        _buttonRegistrations[button] = actionId;

        if (!_sharedHandlers.ContainsKey(actionId))
        {
            System.Action<InputAction.CallbackContext> handler = ctx => OnHotkeyTriggered(actionId);
            _sharedHandlers[actionId] = (handler, action);

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

            action.action.Enable();
        }
    }

    public static void UnregisterHotkey(UXButton button)
    {
        if (button == null || !_buttonRegistrations.TryGetValue(button, out var actionId))
            return;


        if (_hotkeyRegistry.TryGetValue(actionId, out var registrations))
        {
            HotkeyRegistration hotkeyInfo;
            for (int i = registrations.Count - 1; i >= 0; i--)
            {
                if (registrations[i].button == button)
                {
                    hotkeyInfo = registrations[i];
                    registrations.RemoveAt(i);

                    if (registrations.Count == 0 && _sharedHandlers.TryGetValue(actionId, out var handlerInfo))
                    {
                        var (handler, actionRef) = handlerInfo;
                        if (actionRef != null && actionRef.action != null)
                        {
                            actionRef.action.Disable();
                            _sharedHandlers.Remove(actionId);
                            _hotkeyRegistry.Remove(actionId);

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

                    break;
                }
            }
        }


        _buttonRegistrations.Remove(button);
    }

    private static void OnHotkeyTriggered(string actionId)
    {
        if (_hotkeyRegistry.TryGetValue(actionId, out var registrations) && registrations.Count > 0)
        {
            var registration = registrations[registrations.Count - 1];
            registration.button.OnHotkeyTriggered();
        }
    }
}

public static class UXButtonHotkeyExtension
{
    public static void BindHotKey(this UXButton button)
    {
        if (button == null) return;

#if INPUTSYSTEM_SUPPORT
        if (button._hotKeyRefrence != null)
        {
            UXHotkeyComponent.RegisterHotkey(
                button,
                button._hotKeyRefrence,
                button._hotkeyPressType
            );
        }
#endif
    }

    public static void UnBindHotKey(this UXButton button)
    {
        if (button == null) return;
        UXHotkeyComponent.UnregisterHotkey(button);
    }
}
#endif