using System;
using UnityEngine.UIElements;

namespace Uigenx
{
    /// <summary>
    /// Keeps pointer and navigation focus inside a modal. UI Toolkit 6000.5 has no focus scope API,
    /// so the trap is assembled from the parts that do exist: a full screen blocker that swallows
    /// picking, and focus and navigation handlers on the layer root that pull focus back in.
    /// </summary>
    public sealed class UgxModalFocusTrap : IDisposable
    {
        public const string BlockerName = "ugx-modal-blocker";
        public const string BlockerUssClassName = "ugx-modal-blocker";

        readonly VisualElement _layerRoot;
        readonly VisualElement _modalRoot;
        readonly VisualElement _blocker;
        readonly Focusable _previouslyFocused;
        readonly bool _previousDelegatesFocus;

        EventCallback<FocusInEvent> _onFocusIn;
        EventCallback<NavigationMoveEvent> _onNavigationMove;
        bool _disposed;

        UgxModalFocusTrap(VisualElement layerRoot, VisualElement modalRoot, Focusable previouslyFocused)
        {
            _layerRoot = layerRoot;
            _modalRoot = modalRoot;
            _previouslyFocused = previouslyFocused;
            _previousDelegatesFocus = modalRoot.delegatesFocus;

            _blocker = new VisualElement { name = BlockerName, pickingMode = PickingMode.Position, focusable = false };
            _blocker.AddToClassList(BlockerUssClassName);
            _blocker.style.position = Position.Absolute;
            _blocker.style.left = 0f;
            _blocker.style.top = 0f;
            _blocker.style.right = 0f;
            _blocker.style.bottom = 0f;
        }

        public VisualElement Blocker => _blocker;

        public VisualElement ModalRoot => _modalRoot;

        public static UgxModalFocusTrap Install(VisualElement layerRoot, VisualElement modalRoot)
        {
            if (layerRoot == null) throw new ArgumentNullException(nameof(layerRoot));
            if (modalRoot == null) throw new ArgumentNullException(nameof(modalRoot));

            var trap = new UgxModalFocusTrap(layerRoot, modalRoot, layerRoot.focusController?.focusedElement);
            trap.Attach();
            return trap;
        }

        void Attach()
        {
            int index = _layerRoot.IndexOf(_modalRoot);
            if (index < 0) _layerRoot.Add(_blocker);
            else _layerRoot.Insert(index, _blocker);

            _modalRoot.delegatesFocus = true;

            _onFocusIn = OnFocusIn;
            _onNavigationMove = OnNavigationMove;
            _layerRoot.RegisterCallback(_onFocusIn, TrickleDown.TrickleDown);
            _layerRoot.RegisterCallback(_onNavigationMove, TrickleDown.TrickleDown);

            _modalRoot.schedule.Execute(FocusFirst).StartingIn(0);
        }

        public void Dispose()
        {
            if (_disposed) return;
            _disposed = true;

            _layerRoot.UnregisterCallback(_onFocusIn, TrickleDown.TrickleDown);
            _layerRoot.UnregisterCallback(_onNavigationMove, TrickleDown.TrickleDown);
            _onFocusIn = null;
            _onNavigationMove = null;

            _blocker.RemoveFromHierarchy();
            _modalRoot.delegatesFocus = _previousDelegatesFocus;

            if (_previouslyFocused != null && _previouslyFocused.canGrabFocus) _previouslyFocused.Focus();
        }

        void OnFocusIn(FocusInEvent evt)
        {
            if (_disposed) return;
            if (IsInsideModal(evt.target as VisualElement)) return;
            // Focus already moved, so pulling it back has to happen after this event resolves.
            _modalRoot.schedule.Execute(FocusFirst).StartingIn(0);
        }

        void OnNavigationMove(NavigationMoveEvent evt)
        {
            if (_disposed) return;
            if (IsInsideModal(_layerRoot.focusController?.focusedElement as VisualElement)) return;
            evt.StopPropagation();
            FocusFirst();
        }

        bool IsInsideModal(VisualElement element)
        {
            for (var current = element; current != null; current = current.parent)
                if (ReferenceEquals(current, _modalRoot)) return true;
            return false;
        }

        void FocusFirst()
        {
            if (_disposed) return;
            var first = FindFocusable(_modalRoot);
            if (first != null) first.Focus();
            else if (_modalRoot.canGrabFocus) _modalRoot.Focus();
        }

        static VisualElement FindFocusable(VisualElement root)
        {
            if (root.canGrabFocus && root.tabIndex >= 0 && root.enabledInHierarchy) return root;
            for (int i = 0; i < root.childCount; i++)
            {
                var found = FindFocusable(root[i]);
                if (found != null) return found;
            }
            return null;
        }
    }
}