// InputBindingManager.cs

using System;
using System.IO;
using System.Linq;
using System.Collections.Generic;
using System.Reactive.Linq;
using System.Threading.Tasks;
using UnityEngine;
using UnityEngine.InputSystem;
using System.Reactive.Subjects;
using AlicizaX;

using Cysharp.Threading.Tasks;
using RxUnit = System.Reactive.Unit;


    public class InputBindingManager : MonoSingleton<InputBindingManager>
    {
        public const string NULL_BINDING = "__NULL__";
        private const string KEYBOARD_DEVICE = "<Keyboard>";
        private const string MOUSE_DELTA = "<Mouse>/delta";
        private const string MOUSE_SCROLL = "<Mouse>/scroll";
        private const string MOUSE_SCROLL_X = "<Mouse>/scroll/x";
        private const string MOUSE_SCROLL_Y = "<Mouse>/scroll/y";
        private const string KEYBOARD_ESCAPE = "<Keyboard>/escape";

        [Tooltip("InputActionAsset to manage")]
        public InputActionAsset actions;

        public string fileName = "input_bindings.json";
        public bool debugMode = false;

        public Dictionary<string, ActionMap> actionMap = new Dictionary<string, ActionMap>();
        public HashSet<RebindContext> preparedRebinds = new HashSet<RebindContext>();

        internal InputActionRebindingExtensions.RebindingOperation rebindOperation;
        private bool isApplyPending = false;
        private string defaultBindingsJson = string.Empty;
        private string cachedSavePath;
        private Dictionary<string, (ActionMap map, ActionMap.Action action)> actionLookup = new Dictionary<string, (ActionMap, ActionMap.Action)>();

        public readonly ReplaySubject<RxUnit> OnInputsInit = new ReplaySubject<RxUnit>(1);
        public readonly Subject<(bool success, HashSet<RebindContext> appliedContexts)> OnApply = new Subject<(bool, HashSet<RebindContext>)>();
        public readonly Subject<RebindContext> OnRebindPrepare = new Subject<RebindContext>();
        public readonly Subject<RxUnit> OnRebindStart = new Subject<RxUnit>();
        public readonly Subject<(bool success, RebindContext context)> OnRebindEnd = new Subject<(bool, RebindContext)>();
        public readonly Subject<(RebindContext prepared, RebindContext conflict)> OnRebindConflict = new Subject<(RebindContext, RebindContext)>();

        public string SavePath
        {
            get
            {
                if (!string.IsNullOrEmpty(cachedSavePath))
                    return cachedSavePath;

#if UNITY_EDITOR
                string folder = Application.dataPath;
#else
                string folder = Application.persistentDataPath;
#endif
                cachedSavePath = Path.Combine(folder, fileName);
                return cachedSavePath;
            }
        }

        private void EnsureSaveDirectoryExists()
        {
            var directory = Path.GetDirectoryName(SavePath);
            if (!Directory.Exists(directory))
                Directory.CreateDirectory(directory);
        }

        protected override void Awake()
        {
            if (_instance != null && _instance != this)
            {
                Destroy(gameObject);
                return;
            }

            _instance = this;

            if (actions == null)
            {
                Debug.LogError("InputBindingManager: InputActionAsset not assigned.");
                return;
            }

            BuildActionMap();

            try
            {
                defaultBindingsJson = actions.SaveBindingOverridesAsJson();
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[InputBindingManager] Failed to save default bindings: {ex.Message}");
                defaultBindingsJson = string.Empty;
            }

            if (File.Exists(SavePath))
            {
                try
                {
                    var json = File.ReadAllText(SavePath);
                    if (!string.IsNullOrEmpty(json))
                    {
                        actions.LoadBindingOverridesFromJson(json);
                        RefreshBindingPathsFromActions();
                        if (debugMode)
                        {
                            Debug.Log($"Loaded overrides from {SavePath}");
                        }
                    }
                }
                catch (Exception ex)
                {
                    Debug.LogError("Failed to load overrides: " + ex);
                }
            }

            OnInputsInit.OnNext(RxUnit.Default);
            actions.Enable();
        }

        protected override void OnDestroy()
        {
            if (_instance == this)
            {
                _instance = null;
            }

            rebindOperation?.Dispose();
            rebindOperation = null;

            OnInputsInit?.OnCompleted();
            OnApply?.OnCompleted();
            OnRebindPrepare?.OnCompleted();
            OnRebindStart?.OnCompleted();
            OnRebindEnd?.OnCompleted();
            OnRebindConflict?.OnCompleted();
        }

        private void BuildActionMap()
        {
            // Pre-allocate with known capacity to avoid resizing
            int mapCount = actions.actionMaps.Count;
            actionMap.Clear();
            actionLookup.Clear();

            // Estimate total action count for better allocation
            int estimatedActionCount = 0;
            foreach (var map in actions.actionMaps)
            {
                estimatedActionCount += map.actions.Count;
            }

            // Ensure capacity to avoid rehashing
            if (actionMap.Count == 0)
            {
                actionMap = new Dictionary<string, ActionMap>(mapCount);
                actionLookup = new Dictionary<string, (ActionMap, ActionMap.Action)>(estimatedActionCount);
            }

            foreach (var map in actions.actionMaps)
            {
                var actionMapObj = new ActionMap(map);
                actionMap.Add(map.name, actionMapObj);

                // Build lookup dictionary for O(1) action access
                foreach (var actionPair in actionMapObj.actions)
                {
                    actionLookup[actionPair.Key] = (actionMapObj, actionPair.Value);
                }
            }
        }

        private void RefreshBindingPathsFromActions()
        {
            foreach (var mapPair in actionMap.Values)
            {
                foreach (var actionPair in mapPair.actions.Values)
                {
                    var a = actionPair;
                    foreach (var bpair in a.bindings)
                    {
                        bpair.Value.bindingPath.EffectivePath = a.action.bindings[bpair.Key].effectivePath;
                    }
                }
            }
        }

        public sealed class ActionMap
        {
            public string name;
            public Dictionary<string, Action> actions;

            public ActionMap(InputActionMap map)
            {
                name = map.name;
                int actionCount = map.actions.Count;
                actions = new Dictionary<string, Action>(actionCount);
                foreach (var action in map.actions)
                {
                    actions.Add(action.name, new Action(action));
                }
            }

            public sealed class Action
            {
                public InputAction action;
                public Dictionary<int, Binding> bindings;

                public Action(InputAction action)
                {
                    this.action = action;
                    int count = action.bindings.Count;
                    bindings = new Dictionary<int, Binding>(count);

                    for (int i = 0; i < count; i++)
                    {
                        if (action.bindings[i].isComposite)
                        {
                            int first = i + 1;
                            int last = first;
                            while (last < count && action.bindings[last].isPartOfComposite) last++;
                            for (int p = first; p < last; p++)
                                AddBinding(action.bindings[p], p);
                            i = last - 1;
                        }
                        else
                        {
                            AddBinding(action.bindings[i], i);
                        }
                    }

                    void AddBinding(InputBinding binding, int bindingIndex)
                    {
                        bindings.Add(bindingIndex, new Binding(
                            binding.name,
                            action.name,
                            binding.name,
                            bindingIndex,
                            binding.groups?.Split(InputBinding.Separator) ?? Array.Empty<string>(),
                            new BindingPath(binding.path, binding.overridePath),
                            binding
                        ));
                    }
                }

                public readonly struct Binding
                {
                    public readonly string name;
                    public readonly string parentAction;
                    public readonly string compositePart;
                    public readonly int bindingIndex;
                    public readonly string[] group;
                    public readonly BindingPath bindingPath;
                    public readonly InputBinding inputBinding;

                    public Binding(string name, string parentAction, string compositePart, int bindingIndex,
                        string[] group, BindingPath bindingPath, InputBinding inputBinding)
                    {
                        this.name = name;
                        this.parentAction = parentAction;
                        this.compositePart = compositePart;
                        this.bindingIndex = bindingIndex;
                        this.group = group;
                        this.bindingPath = bindingPath;
                        this.inputBinding = inputBinding;
                    }
                }
            }
        }

        public sealed class BindingPath
        {
            public string bindingPath;
            public string overridePath;
            private Subject<RxUnit> observer;

            public BindingPath(string bindingPath, string overridePath)
            {
                this.bindingPath = bindingPath;
                this.overridePath = overridePath;
            }

            public string EffectivePath
            {
                get => !string.IsNullOrEmpty(overridePath) ? overridePath : bindingPath;
                set
                {
                    overridePath = (value == bindingPath) ? string.Empty : value;
                    observer?.OnNext(RxUnit.Default);
                }
            }

            public IObservable<string> EffectivePathObservable
            {
                get
                {
                    observer ??= new Subject<RxUnit>();
                    return observer.Select(_ => EffectivePath);
                }
            }

            public void Dispose()
            {
                observer?.OnCompleted();
                observer?.Dispose();
                observer = null;
            }
        }

        public sealed class RebindContext
        {
            public InputAction action;
            public int bindingIndex;
            public string overridePath;
            private string cachedToString;

            public RebindContext(InputAction action, int bindingIndex, string overridePath)
            {
                this.action = action;
                this.bindingIndex = bindingIndex;
                this.overridePath = overridePath;
            }

            public override bool Equals(object obj)
            {
                if (obj is not RebindContext other) return false;
                if (action == null || other.action == null) return false;
                return action.name == other.action.name && bindingIndex == other.bindingIndex;
            }

            public override int GetHashCode() => (action?.name ?? string.Empty, bindingIndex).GetHashCode();

            public override string ToString()
            {
                if (cachedToString == null && action != null)
                {
                    cachedToString = $"{action.name}:{bindingIndex}";
                }

                return cachedToString ?? "<null>";
            }
        }

        /* ---------------- Public API ---------------- */

        public static InputAction Action(string actionName)
        {
            var instance = Instance;
            if (instance == null) return null;

            if (instance.actionLookup.TryGetValue(actionName, out var result))
            {
                return result.action.action;
            }

            Debug.LogError($"[InputBindingManager] Could not find action '{actionName}'");
            return null;
        }

        public static void StartRebind(string actionName, string compositePartName = null)
        {
            var action = Action(actionName);
            if (action == null) return;

            // decide bindingIndex & deviceMatch automatically
            int bindingIndex = Instance.FindBestBindingIndexForKeyboard(action, compositePartName);
            if (bindingIndex < 0)
            {
                Debug.LogError($"[InputBindingManager] No suitable binding found for action '{actionName}' (part={compositePartName ?? "<null>"})");
                return;
            }

            Instance.actions.Disable();
            Instance.PerformInteractiveRebinding(action, bindingIndex, KEYBOARD_DEVICE, true);
            Instance.OnRebindStart.OnNext(RxUnit.Default);
            if (Instance.debugMode)
            {
                Debug.Log("[InputBindingManager] Rebind started");
            }
        }

        public static void CancelRebind() => Instance.rebindOperation?.Cancel();

        public static async UniTask<bool> ConfirmApply(bool clearConflicts = true)
        {
            if (!Instance.isApplyPending) return false;

            try
            {
                // Create a copy of the prepared rebinds before clearing
                var appliedContexts = new HashSet<RebindContext>(Instance.preparedRebinds);

                foreach (var ctx in Instance.preparedRebinds)
                {
                    if (!string.IsNullOrEmpty(ctx.overridePath))
                    {
                        if (ctx.overridePath == NULL_BINDING)
                        {
                            ctx.action.RemoveBindingOverride(ctx.bindingIndex);
                        }
                        else
                        {
                            ctx.action.ApplyBindingOverride(ctx.bindingIndex, ctx.overridePath);
                        }
                    }

                    var bp = GetBindingPath(ctx.action.name, ctx.bindingIndex);
                    if (bp != null)
                    {
                        bp.EffectivePath = (ctx.overridePath == NULL_BINDING) ? string.Empty : ctx.overridePath;
                    }
                }

                Instance.preparedRebinds.Clear();
                await Instance.WriteOverridesToDiskAsync();
                Instance.OnApply.OnNext((true, appliedContexts));
                Instance.isApplyPending = false;
                if (Instance.debugMode)
                {
                    Debug.Log("[InputBindingManager] Apply confirmed and saved.");
                }

                return true;
            }
            catch (Exception ex)
            {
                Debug.LogError("[InputBindingManager] Failed to apply binds: " + ex);
                Instance.OnApply.OnNext((false, new HashSet<RebindContext>()));
                return false;
            }
        }

        public static void DiscardPrepared()
        {
            if (!Instance.isApplyPending) return;

            // Create a copy of the prepared rebinds before clearing (for event notification)
            var discardedContexts = new HashSet<RebindContext>(Instance.preparedRebinds);

            Instance.preparedRebinds.Clear();
            Instance.isApplyPending = false;
            Instance.OnApply.OnNext((false, discardedContexts));
            if (Instance.debugMode)
            {
                Debug.Log("[InputBindingManager] Prepared rebinds discarded.");
            }
        }

        private void PerformInteractiveRebinding(InputAction action, int bindingIndex, string deviceMatchPath = null, bool excludeMouseMovementAndScroll = true)
        {
            var op = action.PerformInteractiveRebinding(bindingIndex);

            if (!string.IsNullOrEmpty(deviceMatchPath))
            {
                op = op.WithControlsHavingToMatchPath(deviceMatchPath);
            }

            if (excludeMouseMovementAndScroll)
            {
                op = op.WithControlsExcluding(MOUSE_DELTA)
                    .WithControlsExcluding(MOUSE_SCROLL)
                    .WithControlsExcluding(MOUSE_SCROLL_X)
                    .WithControlsExcluding(MOUSE_SCROLL_Y);
            }

            rebindOperation = op
                .OnApplyBinding((o, path) =>
                {
                    if (AnyPreparedRebind(path, action, bindingIndex, out var existing))
                    {
                        PrepareRebind(new RebindContext(action, bindingIndex, path));
                        PrepareRebind(new RebindContext(existing.action, existing.bindingIndex, NULL_BINDING));
                        OnRebindConflict.OnNext((new RebindContext(action, bindingIndex, path), existing));
                    }
                    else if (AnyBindingPath(path, action, bindingIndex, out var dup))
                    {
                        PrepareRebind(new RebindContext(action, bindingIndex, path));
                        PrepareRebind(new RebindContext(dup.action, dup.bindingIndex, NULL_BINDING));
                        OnRebindConflict.OnNext((new RebindContext(action, bindingIndex, path), new RebindContext(dup.action, dup.bindingIndex, dup.action.bindings[dup.bindingIndex].path)));
                    }
                    else
                    {
                        PrepareRebind(new RebindContext(action, bindingIndex, path));
                    }
                })
                .OnComplete(opc =>
                {
                    if (debugMode)
                    {
                        Debug.Log("[InputBindingManager] Rebind completed");
                    }

                    actions.Enable();
                    OnRebindEnd.OnNext((true, new RebindContext(action, bindingIndex, action.bindings[bindingIndex].effectivePath)));
                    CleanRebindOperation();
                })
                .OnCancel(opc =>
                {
                    if (debugMode)
                    {
                        Debug.Log("[InputBindingManager] Rebind cancelled");
                    }

                    actions.Enable();
                    OnRebindEnd.OnNext((false, new RebindContext(action, bindingIndex, action.bindings[bindingIndex].effectivePath)));
                    CleanRebindOperation();
                })
                .WithCancelingThrough(KEYBOARD_ESCAPE)
                .Start();
        }

        private void CleanRebindOperation()
        {
            rebindOperation?.Dispose();
            rebindOperation = null;
        }

        private bool AnyPreparedRebind(string bindingPath, InputAction currentAction, int currentIndex, out RebindContext duplicate)
        {
            foreach (var ctx in preparedRebinds)
            {
                if (ctx.overridePath == bindingPath && (ctx.action != currentAction || (ctx.action == currentAction && ctx.bindingIndex != currentIndex)))
                {
                    duplicate = ctx;
                    return true;
                }
            }

            duplicate = null;
            return false;
        }

        private bool AnyBindingPath(string bindingPath, InputAction currentAction, int currentIndex, out (InputAction action, int bindingIndex) duplicate)
        {
            foreach (var map in actionMap.Values)
            {
                foreach (var actionPair in map.actions.Values)
                {
                    bool isSameAction = actionPair.action == currentAction;

                    foreach (var bindingPair in actionPair.bindings)
                    {
                        // Skip if it's the same action and same binding index
                        if (isSameAction && bindingPair.Key == currentIndex)
                            continue;

                        if (bindingPair.Value.bindingPath.EffectivePath == bindingPath)
                        {
                            duplicate = (actionPair.action, bindingPair.Key);
                            return true;
                        }
                    }
                }
            }

            duplicate = default;
            return false;
        }

        private void PrepareRebind(RebindContext context)
        {
            // Remove existing rebind for same action/binding if exists
            preparedRebinds.Remove(context);

            if (string.IsNullOrEmpty(context.overridePath))
            {
                var bp = GetBindingPath(context.action.name, context.bindingIndex);
                if (bp != null) context.overridePath = bp.bindingPath;
            }

            var bindingPath = GetBindingPath(context.action.name, context.bindingIndex);
            if (bindingPath == null) return;

            if (bindingPath.EffectivePath != context.overridePath)
            {
                preparedRebinds.Add(context);
                isApplyPending = true;
                OnRebindPrepare.OnNext(context);
                if (debugMode)
                {
                    Debug.Log($"Prepared rebind: {context} -> {context.overridePath}");
                }
            }
        }

        private async UniTask WriteOverridesToDiskAsync()
        {
            try
            {
                var json = actions.SaveBindingOverridesAsJson();
                EnsureSaveDirectoryExists();
                using (var sw = new StreamWriter(SavePath, false)) await sw.WriteAsync(json);
                if (debugMode)
                {
                    Debug.Log($"Overrides saved to {SavePath}");
                }
            }
            catch (Exception ex)
            {
                Debug.LogError("Failed to save overrides: " + ex);
                throw;
            }
        }

        public async UniTask ResetToDefaultAsync()
        {
            try
            {
                if (!string.IsNullOrEmpty(defaultBindingsJson))
                {
                    actions.LoadBindingOverridesFromJson(defaultBindingsJson);
                }
                else
                {
                    foreach (var map in actionMap.Values)
                    {
                        foreach (var a in map.actions.Values)
                        {
                            for (int b = 0; b < a.action.bindings.Count; b++)
                            {
                                a.action.RemoveBindingOverride(b);
                            }
                        }
                    }
                }

                RefreshBindingPathsFromActions();
                await WriteOverridesToDiskAsync();
                if (debugMode)
                {
                    Debug.Log("Reset to default and saved.");
                }
            }
            catch (Exception ex)
            {
                Debug.LogError("Failed to reset defaults: " + ex);
            }
        }

        public static BindingPath GetBindingPath(string actionName, int bindingIndex = 0)
        {
            var instance = Instance;
            if (instance == null) return null;

            if (instance.actionLookup.TryGetValue(actionName, out var result))
            {
                if (result.action.bindings.TryGetValue(bindingIndex, out var binding))
                {
                    return binding.bindingPath;
                }
            }

            return null;
        }


        // choose best binding index for keyboard; if compositePartName != null then look for part
        public int FindBestBindingIndexForKeyboard(InputAction action, string compositePartName = null)
        {
            if (action == null) return -1;

            int fallbackPart = -1;
            int fallbackNonComposite = -1;
            bool searchingForCompositePart = !string.IsNullOrEmpty(compositePartName);

            for (int i = 0; i < action.bindings.Count; i++)
            {
                var b = action.bindings[i];

                // If searching for a specific composite part, skip non-matching bindings
                if (searchingForCompositePart)
                {
                    if (!b.isPartOfComposite) continue;
                    if (!string.Equals(b.name, compositePartName, StringComparison.OrdinalIgnoreCase)) continue;
                }

                // Check if this binding is for keyboard
                bool isKeyboardBinding = (!string.IsNullOrEmpty(b.path) && b.path.StartsWith(KEYBOARD_DEVICE, StringComparison.OrdinalIgnoreCase)) ||
                                         (!string.IsNullOrEmpty(b.effectivePath) && b.effectivePath.StartsWith(KEYBOARD_DEVICE, StringComparison.OrdinalIgnoreCase));

                if (b.isPartOfComposite)
                {
                    if (fallbackPart == -1) fallbackPart = i;
                    if (isKeyboardBinding) return i;
                }
                else
                {
                    if (fallbackNonComposite == -1) fallbackNonComposite = i;
                    if (isKeyboardBinding) return i;
                }
            }

            return fallbackNonComposite >= 0 ? fallbackNonComposite : fallbackPart;
        }

        public static InputBindingManager Instance
        {
            get
            {
                if (_instance == null)
                {
                    _instance = FindObjectOfType<InputBindingManager>();
                }

                return _instance;
            }
        }

        private static InputBindingManager _instance;
    }