megakoop/Game/Scripts/Runtime/Pooling/ObjectPooler.cs

using System;
using System.Collections.Generic;
using Game.Scripts.Runtime.Abstractions;
using Game.Scripts.Runtime.Data;
using UnityEngine;
using Unity.Netcode;

namespace Game.Scripts.Runtime.Pooling
{
    [DefaultExecutionOrder(-300)]
    [DisallowMultipleComponent]
    public class ObjectPooler : MonoBehaviour, IPool<EnemyDefinition, GameObject>, IEnemyFactory
    {
        [SerializeField] private Transform poolRoot;
        [SerializeField] private bool markAsDontDestroyOnLoad = true;
        [SerializeField] private int defaultHardCap = 64;

        private readonly Dictionary<string, PoolBucket> _pools = new();
        private readonly Dictionary<GameObject, PoolItem> _lookup = new();

        public static ObjectPooler SharedInstance { get; private set; }

        public event Action<GameObject, EnemyDefinition> InstanceSpawned;
        public event Action<GameObject, EnemyDefinition> InstanceDespawned;

        private void Awake()
        {
            EnsureRoot();

            if (markAsDontDestroyOnLoad && Application.isPlaying)
            {
                DontDestroyOnLoad(gameObject);
            }

            if (SharedInstance == null)
            {
                SharedInstance = this;
            }
        }

        public void Prewarm(EnemyDefinition definition, int count)
        {
            if (definition == null || definition.Prefab == null || count <= 0)
            {
                return;
            }

            var key = GetKey(definition);
            var bucket = GetOrCreateBucket(key);
            bucket.HardCap = bucket.HardCap <= 0 ? defaultHardCap : bucket.HardCap;

            var remaining = Mathf.Max(0, Math.Min(count, bucket.HardCap) - bucket.TotalCount);
            for (var i = 0; i < remaining; i++)
            {
                var instance = CreateInstance(definition, key);
                if (instance == null)
                {
                    return;
                }

                bucket.Available.Enqueue(instance);
            }
        }

        public void SetHardCap(EnemyDefinition definition, int hardCap)
        {
            if (definition == null)
            {
                return;
            }

            var key = GetKey(definition);
            var bucket = GetOrCreateBucket(key);
            bucket.HardCap = Mathf.Max(1, hardCap);
        }

        public GameObject Spawn(EnemyDefinition definition, Vector3 position, Quaternion rotation)
        {
            if (definition == null || definition.Prefab == null)
            {
                Debug.LogWarning("Cannot spawn without a valid definition or prefab.", this);
                return null;
            }

            var key = GetKey(definition);
            var bucket = GetOrCreateBucket(key);
            bucket.HardCap = bucket.HardCap <= 0 ? defaultHardCap : bucket.HardCap;

            if (bucket.Available.Count == 0 && bucket.TotalCount >= bucket.HardCap)
            {
                Debug.LogWarning($"Pool for {key} reached hard cap of {bucket.HardCap}", this);
                return null;
            }

            if (bucket.Available.Count == 0)
            {
                var instance = CreateInstance(definition, key);
                if (instance != null)
                {
                    bucket.Available.Enqueue(instance);
                }
            }

            if (bucket.Available.Count == 0)
            {
                return null;
            }

            var go = bucket.Available.Dequeue();
            if (!_lookup.TryGetValue(go, out var item))
            {
                return null;
            }

            bucket.Active.Add(go);
            SetParentSafely(go, null, true);
            go.transform.SetPositionAndRotation(position + definition.PrefabPivotOffset, rotation);
            go.SetActive(true);

            InstanceSpawned?.Invoke(go, definition);
            return go;
        }

        public void Despawn(GameObject instance)
        {
            if (instance == null)
            {
                return;
            }

            if (!_lookup.TryGetValue(instance, out var item))
            {
                Destroy(instance);
                return;
            }

            var bucket = GetOrCreateBucket(item.Key);
            bucket.Active.Remove(instance);

            instance.SetActive(false);
            var netObj = instance.GetComponent<NetworkObject>();
            SetParentSafely(instance, poolRoot, false, netObj, false);
            if (netObj != null)
            {
                Destroy(netObj);
            }
            bucket.Available.Enqueue(instance);
            item.Handle?.NotifyDespawned();
            InstanceDespawned?.Invoke(instance, item.Definition);
        }

        void IPool<EnemyDefinition, GameObject>.Prewarm(EnemyDefinition key, int count) => Prewarm(key, count);

        GameObject IPool<EnemyDefinition, GameObject>.Get(EnemyDefinition key) => Spawn(key, Vector3.zero, Quaternion.identity);

        void IPool<EnemyDefinition, GameObject>.Release(EnemyDefinition key, GameObject value) => Despawn(value);

        GameObject IEnemyFactory.Spawn(EnemyDefinition definition, Vector3 position, Quaternion rotation) => Spawn(definition, position, rotation);

        void IEnemyFactory.Despawn(GameObject instance) => Despawn(instance);

        private PoolBucket GetOrCreateBucket(string key)
        {
            if (!_pools.TryGetValue(key, out var bucket))
            {
                bucket = new PoolBucket { HardCap = defaultHardCap };
                _pools.Add(key, bucket);
            }

            return bucket;
        }

        private GameObject CreateInstance(EnemyDefinition definition, string key)
        {
            if (definition == null || definition.Prefab == null)
            {
                return null;
            }

            EnsureRoot();
            var go = Instantiate(definition.Prefab);
            SetParentSafely(go, poolRoot, false);
            go.name = $"{definition.Prefab.name}_Pooled";
            go.SetActive(false);

            var handle = go.GetComponent<PooledInstance>();
            if (handle == null)
            {
                handle = go.AddComponent<PooledInstance>();
            }

            handle.Initialize(this, key, definition);
            _lookup[go] = new PoolItem(definition, key, handle);
            return go;
        }

        private static string GetKey(EnemyDefinition definition)
        {
            if (definition == null)
            {
                return string.Empty;
            }

            return string.IsNullOrEmpty(definition.Id) ? definition.name : definition.Id;
        }

        private sealed class PoolBucket
        {
            public readonly Queue<GameObject> Available = new();
            public readonly HashSet<GameObject> Active = new();
            public int HardCap;
            public int TotalCount => Available.Count + Active.Count;
        }

        private sealed class PoolItem
        {
            public readonly EnemyDefinition Definition;
            public readonly string Key;
            public readonly PooledInstance Handle;

            public PoolItem(EnemyDefinition definition, string key, PooledInstance handle)
            {
                Definition = definition;
                Key = key;
                Handle = handle;
            }
        }

        private void EnsureRoot()
        {
            if (poolRoot != null)
            {
                return;
            }

            var root = new GameObject("[Pool]");
            root.transform.SetParent(transform, false);
            poolRoot = root.transform;
        }

        private void SetParentSafely(GameObject instance, Transform parent, bool worldPositionStays, NetworkObject cachedNetworkObject = null, bool restoreAutoSync = true)
        {
            if (instance == null)
            {
                return;
            }

            var netObj = cachedNetworkObject != null ? cachedNetworkObject : instance.GetComponent<NetworkObject>();
            var shouldRestore = false;
            var previousAutoSync = false;

            if (netObj != null && !HasListeningNetworkManager(netObj))
            {
                previousAutoSync = netObj.AutoObjectParentSync;
                netObj.AutoObjectParentSync = false;
                shouldRestore = restoreAutoSync;
            }

            instance.transform.SetParent(parent, worldPositionStays);

            if (shouldRestore && netObj != null)
            {
                netObj.AutoObjectParentSync = previousAutoSync;
            }
        }

        private static bool HasListeningNetworkManager(NetworkObject netObj)
        {
            var manager = netObj != null ? (netObj.NetworkManager != null ? netObj.NetworkManager : NetworkManager.Singleton) : NetworkManager.Singleton;
            return manager != null && manager.IsListening;
        }
    }

    [DisallowMultipleComponent]
    public sealed class PooledInstance : MonoBehaviour
    {
        private ObjectPooler _pool;
        private string _key;
        private EnemyDefinition _definition;

        public event Action<PooledInstance> Despawned;

        internal void Initialize(ObjectPooler pool, string key, EnemyDefinition definition)
        {
            _pool = pool;
            _key = key;
            _definition = definition;
        }

        public void ReturnToPool()
        {
            _pool?.Despawn(gameObject);
        }

        internal void NotifyDespawned()
        {
            Despawned?.Invoke(this);
        }

        public EnemyDefinition Definition => _definition;
        public string Key => _key;
    }
}