Danro/Assets/Resources/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;

        bool wasOpenLastFrame;
        GameObject prevSelectedGlobal;
        readonly List<Selectable> outside = new();
        readonly Dictionary<Selectable, bool> outsidePrev = new();
        Selectable lastModalSelected;

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

        //void OnEnable()
        //{
        //    if (modalRoot != null && modalRoot.activeSelf) OnModalOpened();
        //}

        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)
            {
                if (prevSelectedGlobal && prevSelectedGlobal.activeInHierarchy)
                    EventSystem.current.SetSelectedGameObject(prevSelectedGlobal);
                else
                    EventSystem.current.SetSelectedGameObject(null);
            }

            lastModalSelected = null;
            prevSelectedGlobal = null;
        }

        void CollectOutsideSelectables()
        {
            outside.Clear();
            outsidePrev.Clear();

            // ✅ Новый API Unity 6
            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>();
            }
        }
    }
}