RiftMayhem/Assets/_ProjectRevive/Scripts/ObjectPooling/GameObjectPoolManager.cs

using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Pool;
using UnityEngine.SceneManagement;

/// <summary>
/// Centralized manager for GameObject pooling using Unity's built-in ObjectPool system.
/// Handles per-prefab pools for projectiles, AOE effects, and other frequently instantiated objects.
/// 
/// USAGE:
/// - Replace Object.Instantiate(prefab, pos, rot) with GameObjectPoolManager.Instance.Get(prefab, pos, rot)
/// - Replace Destroy(gameObject) with GameObjectPoolManager.Instance.Release(gameObject)
/// 
/// TO ADD NEW PREFAB TYPES:
/// 1. Add prefab to poolConfigs list in inspector or call RegisterPrefab() at runtime
/// 2. That's it! The system automatically creates pools for any prefab you Get()
/// </summary>
public class GameObjectPoolManager : MonoBehaviour
{
    [System.Serializable]
    public class PoolConfig
    {
        [Tooltip("The prefab to pool")]
        public GameObject prefab;

        [Tooltip("Initial number of objects to create in pool")]
        public int defaultCapacity = 50;

        [Tooltip("Maximum objects in pool before oldest gets destroyed")]
        public int maxSize = 100;

        [Tooltip("Check objects when returned to pool")]
        public bool collectionCheck = true;
    }

    #region Singleton
    private static GameObjectPoolManager _instance;
    public static GameObjectPoolManager Instance
    {
        get
        {
            if (_instance == null)
            {
                _instance = FindFirstObjectByType<GameObjectPoolManager>();
                if (_instance == null)
                {
                    GameObject go = new GameObject("GameObjectPoolManager");
                    _instance = go.AddComponent<GameObjectPoolManager>();
                    DontDestroyOnLoad(go);
                }
            }
            return _instance;
        }
    }

    private void Awake()
    {
        if (_instance == null)
        {
            _instance = this;
            DontDestroyOnLoad(gameObject);

            // Subscribe to scene loaded events
            SceneManager.sceneLoaded += OnSceneLoaded;

            InitializePools();
        }
        else if (_instance != this)
        {
            Destroy(gameObject);
        }
    }

    private void OnDestroy()
    {
        // Unsubscribe from events
        if (_instance == this)
        {
            SceneManager.sceneLoaded -= OnSceneLoaded;
        }
    }
    #endregion

    [Header("Pool Configuration")]
    [SerializeField] private List<PoolConfig> poolConfigs = new List<PoolConfig>();

    [ContextMenu("Reset All Pool Configs")]
    private void ResetPoolConfigs()
    {
        foreach (var config in poolConfigs)
        {
            if (config.prefab != null)
            {
                config.defaultCapacity = 50;
                config.maxSize = 100;
                config.collectionCheck = true;
                Debug.Log($"Reset config for {config.prefab.name}");
            }
        }
        Debug.Log("All pool configs reset to default values");
    }

    [Header("Scene Management")]
    [SerializeField] private bool clearPoolsOnSceneLoad = true;
    [Tooltip("If true, pools will be recreated immediately on scene load. If false, they'll be recreated on-demand.")]
    [SerializeField] private bool preCreatePoolsOnSceneLoad = false;

    [Header("Debug Info")]
    [SerializeField] private bool showDebugLogs = true;
    [SerializeField] private bool showPoolStats = false;

    // Runtime pool storage - maps prefab to its ObjectPool
    private Dictionary<GameObject, ObjectPool<GameObject>> pools = new Dictionary<GameObject, ObjectPool<GameObject>>();

    // Track which pool each active object belongs to for release
    private Dictionary<GameObject, GameObject> activeObjectToPrefab = new Dictionary<GameObject, GameObject>();

    #region Scene Management

    private void OnSceneLoaded(Scene scene, LoadSceneMode mode)
    {
        if (clearPoolsOnSceneLoad)
        {
            if (showDebugLogs)
                Debug.Log($"GameObjectPoolManager: Scene '{scene.name}' loaded, clearing pools");

            ClearAllPools();

            if (preCreatePoolsOnSceneLoad)
            {
                InitializePools();
            }
        }
    }

    /// <summary>
    /// Clears all pools and active object tracking. Called automatically on scene load.
    /// </summary>
    public void ClearAllPools()
    {
        // Clear pools (this will dispose of all pooled objects)
        foreach (var pool in pools.Values)
        {
            pool.Clear();
        }
        pools.Clear();

        // Clear active object tracking
        activeObjectToPrefab.Clear();

        if (showDebugLogs)
            Debug.Log("GameObjectPoolManager: Cleared all pools");
    }

    #endregion

    #region Public API

    /// <summary>
    /// Get a pooled GameObject instance. Creates pool if it doesn't exist.
    /// </summary>
    /// <param name="prefab">Prefab to instantiate</param>
    /// <param name="position">World position</param>
    /// <param name="rotation">World rotation</param>
    /// <returns>Pooled GameObject instance</returns>
    public GameObject Get(GameObject prefab, Vector3 position, Quaternion rotation)
    {
        if (prefab == null)
        {
            Debug.LogError("GameObjectPoolManager: Cannot get null prefab");
            return null;
        }

        // Ensure pool exists
        if (!pools.ContainsKey(prefab))
        {
            CreatePoolForPrefab(prefab);
        }

        // Get object from pool
        GameObject obj = pools[prefab].Get();

        // Double-check the object is valid (in case pool had stale references)
        if (obj == null)
        {
            Debug.LogWarning($"GameObjectPoolManager: Pool returned null object for {prefab.name}. Recreating pool.");
            // Clear and recreate the pool
            pools[prefab].Clear();
            pools.Remove(prefab);
            CreatePoolForPrefab(prefab);
            obj = pools[prefab].Get();
        }

        // Set position and rotation
        obj.transform.position = position;
        obj.transform.rotation = rotation;

        // Track this object
        activeObjectToPrefab[obj] = prefab;

        if (showDebugLogs)
            Debug.Log($"GameObjectPoolManager: Got {prefab.name} from pool");

        return obj;
    }

    /// <summary>
    /// Return a GameObject to its pool.
    /// </summary>
    /// <param name="obj">Object to return to pool</param>
    public void Release(GameObject obj)
    {
        if (obj == null)
        {
            Debug.LogError("GameObjectPoolManager: Cannot release null object");
            return;
        }

        // Find which prefab this object belongs to
        if (!activeObjectToPrefab.TryGetValue(obj, out GameObject prefab))
        {
            Debug.LogWarning($"GameObjectPoolManager: Object {obj.name} not tracked. This might be a non-pooled object. Destroying instead.");
            Destroy(obj);
            return;
        }

        // Check if the pool still exists (might have been cleared on scene change)
        if (!pools.ContainsKey(prefab))
        {
            Debug.LogWarning($"GameObjectPoolManager: Pool for {prefab.name} no longer exists. Destroying object instead.");
            activeObjectToPrefab.Remove(obj);
            Destroy(obj);
            return;
        }

        // Remove from tracking
        activeObjectToPrefab.Remove(obj);

        // Return to pool
        pools[prefab].Release(obj);

        if (showDebugLogs)
            Debug.Log($"GameObjectPoolManager: Released {obj.name} to {prefab.name} pool");
    }

    /// <summary>
    /// Register a new prefab for pooling at runtime.
    /// </summary>
    /// <param name="prefab">Prefab to register</param>
    /// <param name="defaultCapacity">Initial pool size</param>
    /// <param name="maxSize">Maximum pool size</param>
    public void RegisterPrefab(GameObject prefab, int defaultCapacity = 50, int maxSize = 100)
    {
        if (prefab == null) return;

        if (pools.ContainsKey(prefab))
        {
            Debug.LogWarning($"GameObjectPoolManager: Prefab {prefab.name} already registered");
            return;
        }

        PoolConfig config = new PoolConfig
        {
            prefab = prefab,
            defaultCapacity = defaultCapacity,
            maxSize = maxSize,
            collectionCheck = true
        };

        poolConfigs.Add(config);
        CreatePool(config);

        Debug.Log($"GameObjectPoolManager: Registered and created pool for {prefab.name}");
    }

    #endregion

    #region Pool Management

    private void InitializePools()
    {
        foreach (var config in poolConfigs)
        {
            if (config.prefab != null && ValidatePoolConfig(config))
            {
                CreatePool(config);
            }
        }

        Debug.Log($"GameObjectPoolManager: Initialized {pools.Count} pools");
    }

    private bool ValidatePoolConfig(PoolConfig config)
    {
        if (config.defaultCapacity <= 0)
        {
            Debug.LogError($"GameObjectPoolManager: Invalid defaultCapacity ({config.defaultCapacity}) for {config.prefab.name}. Must be > 0. Using default value 50.");
            config.defaultCapacity = 50;
        }

        if (config.maxSize <= 0)
        {
            Debug.LogError($"GameObjectPoolManager: Invalid maxSize ({config.maxSize}) for {config.prefab.name}. Must be > 0. Using default value 100.");
            config.maxSize = 100;
        }

        if (config.maxSize < config.defaultCapacity)
        {
            Debug.LogWarning($"GameObjectPoolManager: maxSize ({config.maxSize}) is less than defaultCapacity ({config.defaultCapacity}) for {config.prefab.name}. Setting maxSize to {config.defaultCapacity}.");
            config.maxSize = config.defaultCapacity;
        }

        return true;
    }

    private void CreatePoolForPrefab(GameObject prefab)
    {
        // Check if we have a config for this prefab
        PoolConfig config = poolConfigs.Find(c => c.prefab == prefab);

        if (config == null)
        {
            // Create default config
            config = new PoolConfig
            {
                prefab = prefab,
                defaultCapacity = 50,
                maxSize = 100,
                collectionCheck = true
            };

            Debug.Log($"GameObjectPoolManager: Auto-creating pool for {prefab.name} with default settings");
        }

        // Always validate config before creating pool
        ValidatePoolConfig(config);
        CreatePool(config);
    }

    private void CreatePool(PoolConfig config)
    {
        if (pools.ContainsKey(config.prefab))
        {
            Debug.LogWarning($"GameObjectPoolManager: Pool for {config.prefab.name} already exists");
            return;
        }

        Debug.Log($"GameObjectPoolManager: Creating pool for {config.prefab.name} with defaultCapacity={config.defaultCapacity}, maxSize={config.maxSize}");

        var pool = new ObjectPool<GameObject>(
            createFunc: () => CreatePooledObject(config.prefab),
            actionOnGet: (obj) => OnGetFromPool(obj),
            actionOnRelease: (obj) => OnReleaseToPool(obj),
            actionOnDestroy: (obj) => OnDestroyPooledObject(obj),
            collectionCheck: config.collectionCheck,
            defaultCapacity: config.defaultCapacity,
            maxSize: config.maxSize
        );

        pools[config.prefab] = pool;

        Debug.Log($"GameObjectPoolManager: Successfully created pool for {config.prefab.name}");
    }

    #endregion

    #region Pool Callbacks

    private GameObject CreatePooledObject(GameObject prefab)
    {
        GameObject obj = Instantiate(prefab);
        obj.name = prefab.name + "(Pooled)";

        // Move to DontDestroyOnLoad so it persists across scenes
        DontDestroyOnLoad(obj);

        // Add PooledObject component to track pool membership
        var pooledComponent = obj.GetComponent<PooledObject>();
        if (pooledComponent == null)
        {
            pooledComponent = obj.AddComponent<PooledObject>();
        }
        pooledComponent.Initialize(prefab);

        if (showDebugLogs)
            Debug.Log($"GameObjectPoolManager: Created new pooled object {obj.name}");

        return obj;
    }

    private void OnGetFromPool(GameObject obj)
    {
        // Check if object is still valid (might be destroyed if scene changed)
        if (obj == null)
        {
            Debug.LogError("GameObjectPoolManager: Pool returned a destroyed object!");
            return;
        }

        // Reset any IPoolable components FIRST
        var poolables = obj.GetComponents<IPoolable>();
        for (int i = 0; i < poolables.Length; i++)
        {
            poolables[i].OnGetFromPool();
        }

        // Activate object AFTER reset
        obj.SetActive(true);
    }

    private void OnReleaseToPool(GameObject obj)
    {
        if (obj == null) return;

        // Reset any IPoolable components
        var poolables = obj.GetComponents<IPoolable>();
        for (int i = 0; i < poolables.Length; i++)
        {
            poolables[i].OnReleaseToPool();
        }

        obj.SetActive(false);
    }

    private void OnDestroyPooledObject(GameObject obj)
    {
        if (obj != null)
        {
            Destroy(obj);
        }
    }

    #endregion

    #region Debug and Stats

    private void OnGUI()
    {
        if (!showPoolStats) return;

        GUILayout.BeginArea(new Rect(10, 10, 300, 300));
        GUILayout.Label("Pool Stats", GUI.skin.box);
        GUILayout.Label($"Active Pools: {pools.Count}");
        GUILayout.Label($"Active Objects: {activeObjectToPrefab.Count}");
        GUILayout.Space(10);

        foreach (var kvp in pools)
        {
            string prefabName = kvp.Key != null ? kvp.Key.name : "NULL";
            int activeCount = 0;
            foreach (var activeKvp in activeObjectToPrefab)
            {
                if (activeKvp.Value == kvp.Key) activeCount++;
            }
            GUILayout.Label($"{prefabName}: {activeCount} active");
        }

        GUILayout.EndArea();
    }

    #endregion
}

/// <summary>
/// Interface for objects that need custom reset behavior when pooled/released.
/// Implement this on components that need to reset state when returned to pool.
/// </summary>
public interface IPoolable
{
    void OnGetFromPool();
    void OnReleaseToPool();
}

/// <summary>
/// Component automatically added to pooled objects to track their source prefab.
/// </summary>
public class PooledObject : MonoBehaviour
{
    [SerializeField] private GameObject sourcePrefab;

    public GameObject SourcePrefab => sourcePrefab;

    public void Initialize(GameObject prefab)
    {
        sourcePrefab = prefab;
    }
}