FlowScope/Packages/com.flowscope.gamecore/Runtime/UI/UIPreloadService.cs

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using FlowScope.Resources;
using UnityEngine;

namespace FlowScope.UI
{
    public sealed class UIPreloadService : IUIPreloadService, IDisposable
    {
        private readonly object _gate = new();
        private readonly IResourceService _resources;
        private readonly Dictionary<Type, IResourceHandle<GameObject>> _handles = new();
        private readonly Dictionary<Type, InFlightPreload> _inFlightLoads = new();
        private int _releaseGeneration;
        private bool _disposed;

        public UIPreloadService(IResourceService resources)
        {
            _resources = resources ?? throw new ArgumentNullException(nameof(resources));
        }

        public async Task PreloadAsync<TPanel>(CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var panelType = typeof(TPanel);
            InFlightPreload inFlightLoad;
            bool startLoad = false;
            string path = null;
            lock (_gate)
            {
                ThrowIfDisposed();

                if (TryGetPreloadedHandle(panelType, out _))
                {
                    return;
                }

                if (!_inFlightLoads.TryGetValue(panelType, out inFlightLoad) ||
                    inFlightLoad.Invalidated ||
                    inFlightLoad.ReleaseGeneration != _releaseGeneration)
                {
                    var attribute = UIManager.GetPanelAttribute(panelType);
                    path = UIManager.ResolvePath(panelType, attribute);
                    inFlightLoad = new InFlightPreload();
                    inFlightLoad.ReleaseGeneration = _releaseGeneration;
                    _inFlightLoads[panelType] = inFlightLoad;
                    startLoad = true;
                }

                inFlightLoad.WaiterCount++;
                if (startLoad)
                {
                    inFlightLoad.Task = LoadAndStoreAsync(panelType, path, inFlightLoad);
                }
            }

            try
            {
                await WaitForSharedLoadAsync(inFlightLoad.Task, cancellationToken);
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                ForgetWaitingCaller(panelType, inFlightLoad);
                throw;
            }
        }

        public bool IsPreloaded<TPanel>()
        {
            lock (_gate)
            {
                if (_disposed)
                {
                    return false;
                }

                return TryGetPreloadedHandle(typeof(TPanel), out _);
            }
        }

        public bool TryGetPreloadedHandle<TPanel>(out IResourceHandle<GameObject> handle)
        {
            lock (_gate)
            {
                if (_disposed)
                {
                    handle = null;
                    return false;
                }

                return TryGetPreloadedHandle(typeof(TPanel), out handle);
            }
        }

        public void Release<TPanel>()
        {
            var panelType = typeof(TPanel);
            IResourceHandle<GameObject> handle = null;
            lock (_gate)
            {
                if (_disposed)
                {
                    return;
                }

                InvalidateInFlightLoad(panelType);
                if (_handles.TryGetValue(panelType, out handle))
                {
                    _handles.Remove(panelType);
                }
            }

            handle?.Dispose();
        }

        public void ReleaseAll()
        {
            List<IResourceHandle<GameObject>> handles;
            lock (_gate)
            {
                _releaseGeneration++;
                handles = new List<IResourceHandle<GameObject>>(_handles.Values);
                _handles.Clear();
            }

            foreach (var handle in handles)
            {
                handle.Dispose();
            }
        }

        public void Dispose()
        {
            lock (_gate)
            {
                if (_disposed)
                {
                    return;
                }

                _disposed = true;
            }

            ReleaseAll();
        }

        private async Task LoadAndStoreAsync(
            Type panelType,
            string path,
            InFlightPreload inFlightLoad)
        {
            IResourceHandle<GameObject> handle = null;
            try
            {
                handle = await _resources.LoadAsync<GameObject>(path, CancellationToken.None);
                if (handle == null || handle.Asset == null)
                {
                    handle?.Dispose();
                    throw new InvalidOperationException($"UI panel resource is missing: {path}");
                }

                var shouldCache = false;
                lock (_gate)
                {
                    ForgetInFlightLoad(panelType, inFlightLoad);
                    if (!_disposed &&
                        !inFlightLoad.Invalidated &&
                        inFlightLoad.ReleaseGeneration == _releaseGeneration &&
                        inFlightLoad.WaiterCount > 0)
                    {
                        _handles[panelType] = handle;
                        shouldCache = true;
                    }
                }

                if (!shouldCache)
                {
                    handle.Dispose();
                }
            }
            catch
            {
                lock (_gate)
                {
                    ForgetInFlightLoad(panelType, inFlightLoad);
                }

                throw;
            }
        }

        private void ForgetInFlightLoad(Type panelType, InFlightPreload inFlightLoad)
        {
            if (_inFlightLoads.TryGetValue(panelType, out var current) &&
                ReferenceEquals(current, inFlightLoad))
            {
                _inFlightLoads.Remove(panelType);
            }
        }

        private void InvalidateInFlightLoad(Type panelType)
        {
            if (_inFlightLoads.TryGetValue(panelType, out var inFlightLoad))
            {
                inFlightLoad.Invalidated = true;
            }
        }

        private void ForgetWaitingCaller(Type panelType, InFlightPreload inFlightLoad)
        {
            lock (_gate)
            {
                if (_inFlightLoads.TryGetValue(panelType, out var current) &&
                    ReferenceEquals(current, inFlightLoad) &&
                    inFlightLoad.WaiterCount > 0)
                {
                    inFlightLoad.WaiterCount--;
                }
            }
        }

        private bool TryGetPreloadedHandle(
            Type panelType,
            out IResourceHandle<GameObject> handle)
        {
            return _handles.TryGetValue(panelType, out handle) &&
                   handle != null &&
                   !handle.IsDisposed &&
                   handle.Asset != null;
        }

        private void ThrowIfDisposed()
        {
            if (_disposed)
            {
                throw new ObjectDisposedException(nameof(UIPreloadService));
            }
        }

        private static async Task WaitForSharedLoadAsync(
            Task loadTask,
            CancellationToken cancellationToken)
        {
            if (!cancellationToken.CanBeCanceled)
            {
                await loadTask;
                return;
            }

            var cancellation = new TaskCompletionSource<bool>(
                TaskCreationOptions.RunContinuationsAsynchronously);
            using (cancellationToken.Register(() => cancellation.TrySetResult(true)))
            {
                if (await Task.WhenAny(loadTask, cancellation.Task) == cancellation.Task)
                {
                    throw new OperationCanceledException(cancellationToken);
                }
            }

            await loadTask;
        }

        private sealed class InFlightPreload
        {
            public Task Task { get; set; }
            public int WaiterCount { get; set; }
            public int ReleaseGeneration { get; set; }
            public bool Invalidated { get; set; }
        }
    }
}