Danro/Assets/Scripts/UI/ModalFocusTrap.cs

using System.Collections.Generic;
using UnityEngine;
using UnityEngine.EventSystems;
using UnityEngine.UI;

namespace RR.Common.UI
{
    [DisallowMultipleComponent]
    public class ModalFocusTrap : MonoBehaviour
    {
        [Header("Modal")]
        public GameObject modalRoot;
        public Selectable firstSelected;

        [Header("Behavior")]
        [Tooltip("Следить и возвращать фокус внутрь модалки, если он уходит наружу.")]
        public bool enforce = true;

        private bool _wasOpenLastFrame;
        private GameObject _prevSelectedGlobal;
        private readonly List<Selectable> _outside = new();
        private readonly Dictionary<Selectable, bool> _outsidePrev = new();
        private Selectable _lastModalSelected;

        void Reset()
        {
            if (modalRoot == null)
            {
                var sibling = transform.parent ? transform.parent.Find("LevelSelectModal") : null;
                if (sibling) modalRoot = sibling.gameObject;
            }
        }

        void Update()
        {
            if (modalRoot == null) return;

            bool isOpen = modalRoot.activeSelf;
            if (isOpen && !_wasOpenLastFrame) OnModalOpened();
            else if (!isOpen && _wasOpenLastFrame) OnModalClosed();

            _wasOpenLastFrame = isOpen;

            if (isOpen && enforce) EnsureFocusInsideModal();
        }

        void OnDisable()
        {
            if (_wasOpenLastFrame) OnModalClosed();
        }

        void OnModalOpened()
        {
            // запомнить, что было выделено до открытия
            _prevSelectedGlobal = EventSystem.current ? EventSystem.current.currentSelectedGameObject : null;

            // выключить интерактив вне модалки
            CollectOutsideSelectables();
            foreach (var s in _outside)
            {
                if (!s) continue;
                _outsidePrev[s] = s.interactable;
                s.interactable = false;
            }

            // сфокусировать элемент внутри модалки
            GameObject toSelect = null;
            if (firstSelected && firstSelected.gameObject.activeInHierarchy)
                toSelect = firstSelected.gameObject;
            else
            {
                var any = modalRoot.GetComponentInChildren<Selectable>(includeInactive: false);
                if (any) toSelect = any.gameObject;
            }

            if (toSelect && EventSystem.current)
            {
                EventSystem.current.SetSelectedGameObject(null);
                EventSystem.current.SetSelectedGameObject(toSelect);
                _lastModalSelected = toSelect.GetComponent<Selectable>();
            }
        }

        void OnModalClosed()
        {
            // Восстанавливаем интерактивность всех внешних элементов
            foreach (var kv in _outsidePrev)
            {
                var s = kv.Key;
                if (s) s.interactable = kv.Value;
            }
            _outsidePrev.Clear();
            _outside.Clear();

            // Возвращаем фокус на элемент, который был выделен до открытия модалки
            if (EventSystem.current)
            {
                GameObject toSelect = null;

                // Проверяем, что предыдущий элемент всё ещё активен и интерактивен
                if (_prevSelectedGlobal && _prevSelectedGlobal.activeInHierarchy)
                {
                    var selectable = _prevSelectedGlobal.GetComponent<Selectable>();
                    if (selectable && selectable.interactable)
                    {
                        toSelect = _prevSelectedGlobal;
                    }
                }

                // Если предыдущий элемент недоступен, ищем любой доступный Selectable
                if (toSelect == null)
                {
                    var anySelectable = Object.FindAnyObjectByType<Selectable>(FindObjectsInactive.Exclude);
                    if (anySelectable && anySelectable.interactable && !IsUnderModal(anySelectable.transform))
                    {
                        toSelect = anySelectable.gameObject;
                    }
                }

                EventSystem.current.SetSelectedGameObject(toSelect);
            }

            _lastModalSelected = null;
            _prevSelectedGlobal = null;
        }

        void CollectOutsideSelectables()
        {
            _outside.Clear();
            _outsidePrev.Clear();

            var all = Object.FindObjectsByType<Selectable>(
                FindObjectsInactive.Include,
                FindObjectsSortMode.None);

            foreach (var s in all)
            {
                if (!s || !s.gameObject.activeInHierarchy) continue;
                if (!IsUnderModal(s.transform))
                    _outside.Add(s);
            }
        }

        bool IsUnderModal(Transform t)
        {
            if (!modalRoot) return false;
            var root = modalRoot.transform;
            for (var cur = t; cur != null; cur = cur.parent)
                if (cur == root) return true;
            return false;
        }

        void EnsureFocusInsideModal()
        {
            if (!EventSystem.current) return;

            var cur = EventSystem.current.currentSelectedGameObject;
            if (cur == null || !IsUnderModal(cur.transform))
            {
                GameObject fallback =
                    (_lastModalSelected && _lastModalSelected.gameObject.activeInHierarchy) ? _lastModalSelected.gameObject :
                    (firstSelected && firstSelected.gameObject.activeInHierarchy) ? firstSelected.gameObject :
                    modalRoot.GetComponentInChildren<Selectable>()?.gameObject;

                if (fallback)
                {
                    EventSystem.current.SetSelectedGameObject(fallback);
                    _lastModalSelected = fallback.GetComponent<Selectable>();
                }
            }
            else
            {
                _lastModalSelected = cur.GetComponent<Selectable>();
            }
        }
    }
}