DeviantMobile-Rohan/Assets/Scripts/Core/ObjectPool.cs

using UnityEngine;
using System.Collections.Generic;

/// <summary>
/// ObjectPool — Generic, lightweight object pool for mobile performance.
/// Eliminates Instantiate/Destroy overhead for frequently spawned objects
/// (VFX, money bags, projectiles, UI elements).
/// 
/// Usage:
///   ObjectPool pool = new ObjectPool(prefab, initialSize: 10);
///   GameObject obj = pool.Get(position, rotation);
///   pool.Return(obj);  // or obj.GetComponent<PooledObject>().ReturnToPool();
/// </summary>
public class ObjectPool
{
    private readonly GameObject _prefab;
    private readonly Transform _container;
    private readonly Queue<GameObject> _available;
    private readonly HashSet<GameObject> _inUse;
    private readonly int _maxSize;
    private readonly bool _allowGrowth;

    public int AvailableCount => _available.Count;
    public int InUseCount => _inUse.Count;
    public int TotalCount => _available.Count + _inUse.Count;

    /// <summary>
    /// Create a new object pool.
    /// </summary>
    /// <param name="prefab">Prefab to pool</param>
    /// <param name="initialSize">Pre-warm count</param>
    /// <param name="maxSize">Hard cap (0 = unlimited)</param>
    /// <param name="parent">Optional parent transform for organization</param>
    /// <param name="allowGrowth">Allow creating beyond initialSize if pool is empty</param>
    public ObjectPool(GameObject prefab, int initialSize = 10, int maxSize = 0, Transform parent = null, bool allowGrowth = true)
    {
        _prefab = prefab;
        _maxSize = maxSize;
        _allowGrowth = allowGrowth;
        _available = new Queue<GameObject>(initialSize);
        _inUse = new HashSet<GameObject>();

        // Create container for organization
        if (parent == null)
        {
            GameObject containerObj = new GameObject($"Pool_{prefab.name}");
            _container = containerObj.transform;
            Object.DontDestroyOnLoad(containerObj);
        }
        else
        {
            _container = parent;
        }

        // Pre-warm
        for (int i = 0; i < initialSize; i++)
        {
            GameObject obj = CreateNewInstance();
            obj.SetActive(false);
            _available.Enqueue(obj);
        }
    }

    /// <summary>
    /// Get an object from the pool. Activates it at the given position/rotation.
    /// </summary>
    public GameObject Get(Vector3 position, Quaternion rotation)
    {
        GameObject obj = null;

        // Try to get from pool
        while (_available.Count > 0)
        {
            obj = _available.Dequeue();
            if (obj != null) break; // Found valid object
            obj = null; // Was destroyed externally, skip
        }

        // Pool empty — create new if allowed
        if (obj == null)
        {
            if (!_allowGrowth && _maxSize > 0 && TotalCount >= _maxSize)
                return null; // Hard cap reached

            obj = CreateNewInstance();
        }

        obj.transform.SetPositionAndRotation(position, rotation);
        obj.SetActive(true);
        _inUse.Add(obj);

        // Notify pooled components
        var pooled = obj.GetComponent<PooledObject>();
        if (pooled != null) pooled.OnGetFromPool();

        return obj;
    }

    /// <summary>
    /// Return an object to the pool. Deactivates it.
    /// </summary>
    public void Return(GameObject obj)
    {
        if (obj == null) return;

        _inUse.Remove(obj);

        // Notify pooled components
        var pooled = obj.GetComponent<PooledObject>();
        if (pooled != null) pooled.OnReturnToPool();

        obj.SetActive(false);
        obj.transform.SetParent(_container);
        _available.Enqueue(obj);
    }

    /// <summary>
    /// Return all in-use objects to the pool.
    /// </summary>
    public void ReturnAll()
    {
        // Copy to avoid modification during iteration
        var inUseList = new List<GameObject>(_inUse);
        foreach (var obj in inUseList)
        {
            if (obj != null) Return(obj);
        }
    }

    /// <summary>
    /// Destroy all pooled objects and clear the pool.
    /// </summary>
    public void Dispose()
    {
        foreach (var obj in _available)
            if (obj != null) Object.Destroy(obj);

        foreach (var obj in _inUse)
            if (obj != null) Object.Destroy(obj);

        _available.Clear();
        _inUse.Clear();

        if (_container != null)
            Object.Destroy(_container.gameObject);
    }

    private GameObject CreateNewInstance()
    {
        GameObject obj = Object.Instantiate(_prefab, _container);
        obj.name = _prefab.name; // Remove "(Clone)" suffix

        // Add PooledObject component if not present
        var pooled = obj.GetComponent<PooledObject>();
        if (pooled == null)
        {
            pooled = obj.AddComponent<PooledObject>();
        }
        pooled.OwnerPool = this;

        return obj;
    }
}

/// <summary>
/// PooledObject — Attach to pooled prefabs. Provides auto-return and lifecycle callbacks.
/// </summary>
public class PooledObject : MonoBehaviour
{
    public ObjectPool OwnerPool { get; set; }

    [Tooltip("Auto-return to pool after this many seconds (0 = manual return only)")]
    [SerializeField] private float autoReturnDelay = 0f;

    private float _activeTimer;
    private Coroutine _autoReturnCoroutine;

    /// <summary>Called when taken from pool.</summary>
    public virtual void OnGetFromPool()
    {
        _activeTimer = 0f;
        // Start coroutine-based auto-return instead of per-frame Update check
        if (autoReturnDelay > 0f)
        {
            if (_autoReturnCoroutine != null) StopCoroutine(_autoReturnCoroutine);
            _autoReturnCoroutine = StartCoroutine(AutoReturnAfterDelay());
        }
    }

    /// <summary>Called when returned to pool.</summary>
    public virtual void OnReturnToPool()
    {
        _activeTimer = 0f;
        if (_autoReturnCoroutine != null)
        {
            StopCoroutine(_autoReturnCoroutine);
            _autoReturnCoroutine = null;
        }
    }

    /// <summary>Return this object to its pool.</summary>
    public void ReturnToPool()
    {
        if (_autoReturnCoroutine != null)
        {
            StopCoroutine(_autoReturnCoroutine);
            _autoReturnCoroutine = null;
        }
        if (OwnerPool != null)
            OwnerPool.Return(gameObject);
        else
            gameObject.SetActive(false);
    }

    /// <summary>Coroutine-based auto-return — eliminates per-frame Update overhead.</summary>
    private System.Collections.IEnumerator AutoReturnAfterDelay()
    {
        yield return new WaitForSeconds(autoReturnDelay);
        ReturnToPool();
    }
}

/// <summary>
/// PoolManager — Centralized pool registry. Access pools by prefab reference.
/// Singleton, persists across scenes.
/// </summary>
public class PoolManager : MonoBehaviour
{
    private static PoolManager _instance;
    public static PoolManager Instance
    {
        get
        {
            if (_instance == null)
            {
                var go = new GameObject("PoolManager");
                _instance = go.AddComponent<PoolManager>();
                DontDestroyOnLoad(go);
            }
            return _instance;
        }
    }

    private readonly Dictionary<GameObject, ObjectPool> _pools = new Dictionary<GameObject, ObjectPool>();

    /// <summary>
    /// Get or create a pool for a prefab.
    /// </summary>
    public ObjectPool GetPool(GameObject prefab, int initialSize = 10, int maxSize = 0)
    {
        if (_pools.TryGetValue(prefab, out ObjectPool pool))
            return pool;

        pool = new ObjectPool(prefab, initialSize, maxSize, transform);
        _pools[prefab] = pool;
        return pool;
    }

    /// <summary>
    /// Shorthand: Get an object from a prefab's pool.
    /// </summary>
    public GameObject Spawn(GameObject prefab, Vector3 position, Quaternion rotation, int poolSize = 10)
    {
        return GetPool(prefab, poolSize).Get(position, rotation);
    }

    /// <summary>
    /// Shorthand: Return an object to its pool.
    /// </summary>
    public void Despawn(GameObject obj)
    {
        var pooled = obj.GetComponent<PooledObject>();
        if (pooled != null && pooled.OwnerPool != null)
            pooled.OwnerPool.Return(obj);
        else
            obj.SetActive(false);
    }

    /// <summary>
    /// Pre-warm a pool for a prefab.
    /// </summary>
    public void PreWarm(GameObject prefab, int count)
    {
        GetPool(prefab, count);
    }

    /// <summary>
    /// Clear all pools (call on scene transitions).
    /// </summary>
    public void ClearAll()
    {
        foreach (var pool in _pools.Values)
            pool.Dispose();
        _pools.Clear();
    }

    private void OnDestroy()
    {
        ClearAll();
        _instance = null;
    }
}