Danro/Assets/Scripts/UI/ModalFocusTrap.cs

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

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;

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

        private void Update()
        {
            if (!modalRoot) return;

            bool isOpen = modalRoot.activeSelf;
            switch (isOpen)
            {
                case true when !_wasOpenLastFrame:
                    OnModalOpened();
                    break;
                case false when _wasOpenLastFrame:
                    OnModalClosed();
                    break;
            }

            _wasOpenLastFrame = isOpen;

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

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

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

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

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

            if (!toSelect || !EventSystem.current) return;
            EventSystem.current.SetSelectedGameObject(null);
            EventSystem.current.SetSelectedGameObject(toSelect);
            _lastModalSelected = toSelect.GetComponent<Selectable>();
        }

        private void OnModalClosed()
        {
            // Восстанавливаем интерактивность всех внешних элементов
            foreach ((Selectable selectable, bool value) in _outsidePrev)
            {
                if (selectable) selectable.interactable = value;
            }
            _outsidePrev.Clear();
            _outside.Clear();

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

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

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

                EventSystem.current.SetSelectedGameObject(toSelect);
            }

            _lastModalSelected = null;
            _prevSelectedGlobal = null;
        }

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

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

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

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

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

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

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