#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 InputActionReference reference;
        public readonly EHotkeyPressType pressType;
        public readonly UXButton button;

        public HotkeyRegistration(UXButton btn, InputActionReference reference, EHotkeyPressType pressType)
        {
            button = btn;
            this.reference = reference;
            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, EHotkeyPressType pressType)> _sharedHandlers =
        new Dictionary<string, (System.Action<InputAction.CallbackContext>, EHotkeyPressType)>(32);


    private static readonly Dictionary<UXButton, HashSet<string>> _buttonRegistrations =
        new Dictionary<UXButton, HashSet<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, action, pressType);


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

        registrations.Add(registration);


        if (!_buttonRegistrations.TryGetValue(button, out var actionIds))
        {
            actionIds = new HashSet<string>();
            _buttonRegistrations[button] = actionIds;
        }

        actionIds.Add(actionId);

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

            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 actionIds))
            return;

        foreach (var actionId in actionIds)
        {
            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];
                        if (_sharedHandlers.TryGetValue(actionId, out var handlerInfo))
                        {
                            var (handler, pressType) = handlerInfo;
                            var actionRef = hotkeyInfo.reference;

                            if (actionRef != null && actionRef.action != null)
                            {
                                if (registrations.Count == 0)
                                {
                                    actionRef.action.Disable();
                                }

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

                            _sharedHandlers.Remove(actionId);
                        }

                        registrations.RemoveAt(i);
                        _hotkeyRegistry.Remove(actionId);
                        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