canto/Assets/Scripts/UI/Orchestration/UiSceneManager.cs

using UnityEngine;
using R3;
using System.Collections.Generic;
using System;
using UnityEngine.UIElements;
using UnityEngine.Pool;
using KitsuneCafe.System;

namespace KitsuneCafe.UI
{
    public class UISceneManager : MonoBehaviour
    {
        [SerializeField] private UiOrchestrator orchestrator;

        [SerializeField] private Transform root;
        [SerializeField] private PanelSettings panelSettings;
        [SerializeField] private SerializableDuration despawnTimeout;
        [SerializeField] private GameObject prefab;

        private readonly Dictionary<ElementId, GameObject> activeElements = new();
        private readonly Dictionary<VisualTreeAsset, IObjectPool<GameObject>> pools = new();

        private void Awake()
        {
            var d = Disposable.CreateBuilder();

            Observable.FromEvent<EventHandler<SpawnElementRequest>, SpawnElementRequest>(
                h => (sender, e) => h(e),
                e => orchestrator.SpawnRequested += e,
                e => orchestrator.SpawnRequested -= e
            ).Subscribe(request =>
            {
                SpawnElement(request.Id, request.Element, request.Position);
            })
            .AddTo(ref d);

            Observable.FromEvent<EventHandler<DespawnElementRequest>, DespawnElementRequest>(
                h => (sender, e) => h(e),
                e => orchestrator.DespawnRequested += e,
                e => orchestrator.DespawnRequested -= e
            )
            .Select(request => request.Id)
            .Subscribe(DespawnElement)
            .AddTo(ref d);

            d.RegisterTo(destroyCancellationToken);
        }

        private IObjectPool<GameObject> GetOrCreatePool(IUiElement element)
        {
            if (!pools.TryGetValue(element.VisualTreeAsset, out var pool))
            {
                pool = new ObjectPool<GameObject>(
                    () => CreatePoolObject(pool),
                    GetPoolObject,
                    ReleasePoolObject,
                    DestroyPoolObject,
                    collectionCheck: false,
                    defaultCapacity: 10,
                    maxSize: 100
                );

                pools.Add(element.VisualTreeAsset, pool);
            }
            return pool;
        }

        private GameObject CreatePoolObject(IObjectPool<GameObject> pool)
        {
            GameObject obj;
            if (prefab != null)
            {
                obj = Instantiate(prefab);
            }
            else
            {
                obj = new GameObject(
                    "Pooled Object",
                    typeof(UIDocument),
                    typeof(PooledObject),
                    typeof(UiElementInstance)
                );
            }

            var doc = obj.GetComponent<UIDocument>();
            doc.panelSettings = panelSettings;

            var poolObj = obj.GetComponent<PooledObject>();
            poolObj.objectPool = pool;

            var instance = obj.GetComponent<UiElementInstance>();
            instance.Document = doc;
            instance.PooledObject = poolObj;

            obj.transform.SetParent(root);

            return obj;
        }

        private void GetPoolObject(GameObject obj)
        {
            obj.SetActive(true);
        }

        private void ReleasePoolObject(GameObject obj)
        {
            obj.SetActive(false);
        }

        private void DestroyPoolObject(GameObject obj)
        {
            Destroy(obj);
        }

        private GameObject GetUiDocument(IUiElement element)
        {
            return GetOrCreatePool(element).Get();
        }

        private void SpawnElement(ElementId id, IUiElement element, Vector3 worldPosition)
        {
            var obj = GetUiDocument(element);
            activeElements[id] = obj;

            var trans = obj.transform;
            trans.SetParent(root);
            trans.SetPositionAndRotation(worldPosition, Quaternion.identity);
            trans.localScale = Vector3.one;

            if (obj.TryGetComponent(out UiElementInstance instance))
            {
                instance.Configure(element, panelSettings);
            }
        }

        // FIXME: track which pool each active element belongs to?
        private void DespawnElement(ElementId id)
        {
            if (activeElements.TryGetValue(id, out GameObject obj))
            {
                if (obj.TryGetComponent<UiElementInstance>(out var instance))
                {
                    instance.Despawn()
                        .Race(Observable.Timer(despawnTimeout).Do(_ => instance.DespawnNow()))
                        .Do(onCompleted: _ => activeElements.Remove(id))
                        .Subscribe()
                        .AddTo(obj);

                }
                else if (obj.TryGetComponent(out PooledObject poolObj))
                {
                    poolObj.Release();
                    activeElements.Remove(id);
                }
            }
        }
    }
}