AppleHillsProduction/Assets/Scripts/Pooling/MultiPrefabPool.cs

using System.Collections.Generic;
using UnityEngine;

namespace Pooling
{
    /// <summary>
    /// Object pool that supports multiple prefab types with usage tracking and intelligent trimming.
    /// </summary>
    /// <typeparam name="T">The type of component to pool</typeparam>
    public abstract class MultiPrefabPool<T> : MonoBehaviour where T : Component
    {
        [Tooltip("Whether to pre-instantiate objects during initialization or create them on demand")]
        public bool preInstantiatePrefabs = false;
        
        [Tooltip("Initial number of objects to pre-instantiate per prefab (if preInstantiatePrefabs is true)")]
        public int initialObjectsPerPrefab = 2;
        
        [Tooltip("Maximum number of objects to keep in the pool across all prefab types")]
        public int totalMaxPoolSize = 50;
        
        [Tooltip("Maximum number of inactive instances to keep per prefab type")]
        public int maxPerPrefabPoolSize = 5;
        
        protected Dictionary<int, Stack<T>> pooledObjects = new Dictionary<int, Stack<T>>();
        protected Dictionary<int, int> prefabUsageCount = new Dictionary<int, int>();
        protected List<T> prefabs;
        protected int totalPooledCount = 0;
        
        /// <summary>
        /// Initialize the pool with the specified prefabs
        /// </summary>
        /// <param name="prefabsToPool">The list of prefabs to use for this pool</param>
        public virtual void Initialize(List<T> prefabsToPool)
        {
            prefabs = prefabsToPool;
            
            // Initialize usage tracking
            for (int i = 0; i < prefabs.Count; i++)
            {
                prefabUsageCount[i] = 0;
                pooledObjects[i] = new Stack<T>();
            }
            
            // Pre-instantiate objects only if enabled
            if (preInstantiatePrefabs)
            {
                // Calculate how many to pre-instantiate based on available pool size
                int totalToCreate = Mathf.Min(totalMaxPoolSize, prefabs.Count * initialObjectsPerPrefab);
                int perPrefab = Mathf.Max(1, totalToCreate / prefabs.Count);
                
                for (int i = 0; i < prefabs.Count; i++)
                {
                    for (int j = 0; j < perPrefab; j++)
                    {
                        if (totalPooledCount >= totalMaxPoolSize) break;
                        CreateNew(i);
                    }
                }
            }
            
            Debug.Log($"[{GetType().Name}] Initialized with {prefabs.Count} prefab types");
        }
        
        /// <summary>
        /// Creates a new instance and adds it to the pool
        /// </summary>
        protected virtual T CreateNew(int prefabIndex)
        {
            if (prefabs == null || prefabIndex >= prefabs.Count)
            {
                Debug.LogError($"[{GetType().Name}] Invalid prefab index or prefabs list is null!");
                return null;
            }
            
            T prefab = prefabs[prefabIndex];
            T obj = Instantiate(prefab, transform);
            obj.gameObject.SetActive(false);
            
            // Initialize IPoolable components if present
            IPoolable poolable = obj.GetComponent<IPoolable>();
            if (poolable != null)
            {
                poolable.OnDespawn();
            }
            
            if (!pooledObjects.ContainsKey(prefabIndex))
            {
                pooledObjects[prefabIndex] = new Stack<T>();
            }
            
            pooledObjects[prefabIndex].Push(obj);
            totalPooledCount++;
            return obj;
        }
        
        /// <summary>
        /// Gets an object from the pool for the specified prefab index, or creates a new one if the pool is empty
        /// </summary>
        /// <param name="prefabIndex">The index of the prefab to get from the pool</param>
        /// <returns>An object ready to use</returns>
        public virtual T Get(int prefabIndex)
        {
            T obj;
            
            // Track usage frequency
            if (prefabUsageCount.ContainsKey(prefabIndex))
            {
                prefabUsageCount[prefabIndex]++;
            }
            else
            {
                prefabUsageCount[prefabIndex] = 1;
            }
            
            if (pooledObjects.ContainsKey(prefabIndex) && pooledObjects[prefabIndex].Count > 0)
            {
                obj = pooledObjects[prefabIndex].Pop();
                totalPooledCount--;
            }
            else
            {
                // Create new object without adding to pool
                T prefab = prefabs[prefabIndex];
                obj = Instantiate(prefab, transform);
            }
            
            obj.gameObject.SetActive(true);
            
            // Call OnSpawn for IPoolable components
            IPoolable poolable = obj.GetComponent<IPoolable>();
            if (poolable != null)
            {
                poolable.OnSpawn();
            }
            
            return obj;
        }
        
        /// <summary>
        /// Returns an object to the pool
        /// </summary>
        /// <param name="obj">The object to return</param>
        /// <param name="prefabIndex">The index of the prefab this object was created from</param>
        public virtual void Return(T obj, int prefabIndex)
        {
            if (obj == null) return;
            
            // Call OnDespawn for IPoolable components
            IPoolable poolable = obj.GetComponent<IPoolable>();
            if (poolable != null)
            {
                poolable.OnDespawn();
            }
            
            // Check if we're under the maximum pool size for this prefab type
            bool keepObject = totalPooledCount < totalMaxPoolSize;
            
            // Additional constraint: don't keep too many of any single prefab type
            if (pooledObjects.ContainsKey(prefabIndex) && 
                pooledObjects[prefabIndex].Count >= maxPerPrefabPoolSize)
            {
                keepObject = false;
            }
            
            if (keepObject)
            {
                obj.gameObject.SetActive(false);
                obj.transform.SetParent(transform);
                
                if (!pooledObjects.ContainsKey(prefabIndex))
                {
                    pooledObjects[prefabIndex] = new Stack<T>();
                }
                
                pooledObjects[prefabIndex].Push(obj);
                totalPooledCount++;
            }
            else
            {
                Destroy(obj.gameObject);
            }
        }
        
        /// <summary>
        /// Trims the pool to remove excess objects
        /// Can be called periodically or when memory pressure is high
        /// </summary>
        public virtual void TrimExcess()
        {
            // If we're under the limit, no need to trim
            if (totalPooledCount <= totalMaxPoolSize) return;
            
            // Calculate how many to remove
            int excessCount = totalPooledCount - totalMaxPoolSize;
            
            // Get prefab indices sorted by usage (least used first)
            List<KeyValuePair<int, int>> sortedUsage = new List<KeyValuePair<int, int>>(prefabUsageCount);
            sortedUsage.Sort((a, b) => a.Value.CompareTo(b.Value));
            
            // Remove objects from least used prefabs first
            foreach (var usage in sortedUsage)
            {
                int prefabIndex = usage.Key;
                if (!pooledObjects.ContainsKey(prefabIndex) || pooledObjects[prefabIndex].Count == 0) continue;
                
                // How many to remove from this prefab type
                int toRemove = Mathf.Min(pooledObjects[prefabIndex].Count, excessCount);
                
                for (int i = 0; i < toRemove; i++)
                {
                    if (pooledObjects[prefabIndex].Count == 0) break;
                    
                    T obj = pooledObjects[prefabIndex].Pop();
                    Destroy(obj.gameObject);
                    totalPooledCount--;
                    excessCount--;
                    
                    if (excessCount <= 0) return;
                }
            }
        }
        
        /// <summary>
        /// Logs pool statistics to the console
        /// </summary>
        public virtual void LogPoolStats()
        {
            Debug.Log($"[{GetType().Name}] Total pooled objects: {totalPooledCount}/{totalMaxPoolSize}");
            
            string prefabDetails = "";
            int index = 0;
            foreach (var entry in pooledObjects)
            {
                int prefabIndex = entry.Key;
                int count = entry.Value.Count;
                int usageCount = prefabUsageCount.ContainsKey(prefabIndex) ? prefabUsageCount[prefabIndex] : 0;
                
                string prefabName = prefabIndex < prefabs.Count ? prefabs[prefabIndex].name : "Unknown";
                prefabDetails += $"\n  - {prefabName}: {count} pooled, {usageCount} usages";
                
                // Limit the output to avoid too much text
                if (++index >= 10 && pooledObjects.Count > 10)
                {
                    prefabDetails += $"\n  - ...and {pooledObjects.Count - 10} more prefab types";
                    break;
                }
            }
            
            Debug.Log($"[{GetType().Name}] Pool details:{prefabDetails}");
        }
        
#if UNITY_EDITOR
        private float _lastLogTime = 0f;
        
        protected virtual void Update()
        {
            // Log pool stats every 5 seconds if in the editor
            if (Time.time - _lastLogTime > 5f)
            {
                LogPoolStats();
                _lastLogTime = Time.time;
            }
        }
#endif
    }
}