XRLib/Assets/Scripts/Simulator/Components/ComponentPoolManager.cs

using Cysharp.Threading.Tasks;
using System;
using System.Collections.Generic;
using UnityEngine;
using UVC.Data;
using UVC.Data.Core;
using UVC.Data.Mqtt;
using UVC.Pool;
using UVC.UI.Loading;

namespace Simulator.Data
{
    public class ComponentPoolManager
    {
        private readonly Transform _parentTransform;

        private readonly Dictionary<string, string> _prefabPaths = new Dictionary<string, string>()
        {
            {"source","prefabs/pallet"},
            {"queue","prefabs/queue"},
            {"rack","prefabs/rack"},
            {"sink","prefabs/Sink_Container"},
            {"asrs","prefabs/ASRS"},
            {"robotArm","prefabs/RobotArm"},
            {"processor","prefabs/Processor"},
            {"worker","prefabs/Worker"}
        };

        private GameObjectPool<SourceComponent> _sourcePool;
        private GameObjectPool<SinkComponent> _sinkPool;
        private GameObjectPool<QueueComponent> _queuePool;
        private GameObjectPool<RackComponent> _rackPool;
        private GameObjectPool<ASRSComponent> _asrsPool;
        private GameObjectPool<RobotArmComponent> _robotArmPool;
        private GameObjectPool<ProcessorComponent> _processorPool;
        private GameObjectPool<WorkerComponent> _workerPool;

        private readonly Dictionary<ComponentType, DataMapper> _dataMapperDict = new Dictionary<ComponentType, DataMapper>();

        public IReadOnlyDictionary<ComponentType, DataMapper> DataMappers => _dataMapperDict;

        private ComponentUIBinder _uiBinder;
        public ComponentUIBinder UIBinder => _uiBinder;

        public ComponentPoolManager(Transform parentTransform)
        {
            _parentTransform = parentTransform;
        }

        public void InitializeUIBinder()
        {
            _uiBinder = new ComponentUIBinder();
            _uiBinder.Initialize();
        }

        public void InitializeAllPools()
        {
            InitializePoolAsync<SourceComponent>(_sourcePool, "source").ContinueWith(pool => { _sourcePool = pool; _dataMapperDict.Add(ComponentType.Source, CreateDataMapper(ComponentType.Source)); });
            InitializePoolAsync<SinkComponent>(_sinkPool, "sink").ContinueWith(pool => { _sinkPool = pool; _dataMapperDict.Add(ComponentType.Sink, CreateDataMapper(ComponentType.Sink)); });
            InitializePoolAsync<QueueComponent>(_queuePool, "queue").ContinueWith(pool => { _queuePool = pool; _dataMapperDict.Add(ComponentType.Queue, null); });
            InitializePoolAsync<RackComponent>(_rackPool, "rack").ContinueWith(pool => { _rackPool = pool; _dataMapperDict.Add(ComponentType.Rack, CreateDataMapper(ComponentType.Rack)); });
            InitializePoolAsync<ASRSComponent>(_asrsPool, "asrs").ContinueWith(pool => { _asrsPool = pool; _dataMapperDict.Add(ComponentType.ASRS, CreateDataMapper(ComponentType.ASRS)); });
            InitializePoolAsync<RobotArmComponent>(_robotArmPool, "robotArm").ContinueWith(pool => { _robotArmPool = pool; _dataMapperDict.Add(ComponentType.RobotArm, CreateDataMapper(ComponentType.RobotArm)); });
            InitializePoolAsync<ProcessorComponent>(_processorPool, "processor").ContinueWith(pool => { _processorPool = pool; _dataMapperDict.Add(ComponentType.Processor, CreateDataMapper(ComponentType.Processor)); });
            InitializePoolAsync<WorkerComponent>(_workerPool, "worker").ContinueWith(pool => { _workerPool = pool; _dataMapperDict.Add(ComponentType.Worker, CreateDataMapper(ComponentType.Worker)); });
        }

        private DataMapper CreateDataMapper(ComponentType type)
        {
            switch (type)
            {
                case ComponentType.Source:
                    return new DataMapper(CreateEventMask("Source", new DataMask
                    {
                        ["entity_ids"] = new List<string>(),
                        ["total_entity"] = 0
                    }));
                case ComponentType.Sink:
                    return new DataMapper(CreateEventMask("Sink", new DataMask
                    {
                        ["entity_ids"] = new List<string>()
                    }));
                case ComponentType.RobotArm:
                    return new DataMapper(CreateEventMask("RobotArm", new DataMask
                    {
                        ["entity_id"] = "",
                        ["processing_time"] = 0.0f,
                        ["from"] = "",
                        ["to"] = ""
                    }));
                case ComponentType.Rack:
                    return new DataMapper(new DataMask { ObjectName = "rack", ObjectIdKey = "component_id" });
                case ComponentType.ASRS:
                    return new DataMapper(new DataMask { ObjectName = "asrs", ObjectIdKey = "component_id" });
                default:
                    return new DataMapper(new DataMask());
            }
        }

        private DataMask CreateEventMask(string objectName, DataMask dataMask)
        {
            var mask = new DataMask
            {
                ObjectName = objectName,
                ObjectIdKey = "component_id",
                ["component_id"] = "",
                ["event_name"] = "",
                ["timestamp"] = new DateTime(),
                ["data"] = dataMask
            };
            return mask;
        }

        #region Pool Initialization

        private async UniTask<GameObjectPool<T>> InitializePoolAsync<T>(GameObjectPool<T> existingPool, string key) where T : MonoBehaviour
        {
            if (existingPool != null) return existingPool;
            var prefab = await Resources.LoadAsync<GameObject>(_prefabPaths[key]) as GameObject;
            if (prefab == null) { Debug.LogError($"Prefab not found at path: {_prefabPaths[key]}"); return null; }
            return new GameObjectPool<T>(prefab, _parentTransform);
        }

        #endregion

        public Dictionary<string, ComponentBase> SpawnComponents(LogicDetailData data)
        {
            var registry = new Dictionary<string, ComponentBase>();

            var ps = data.production_system;
            if (ps != null)
            {
                SpawnGroup(ps.sources, _sourcePool, ComponentType.Source, (c, d) => c.SetComponent(d), registry);
                SpawnGroup(ps.sinks, _sinkPool, ComponentType.Sink, (c, d) => c.SetComponent(d), registry);
                SpawnGroup(ps.racks, _rackPool, ComponentType.Rack, (c, d) => c.SetComponent(d), registry);
                SpawnGroup(ps.asrs, _asrsPool, ComponentType.ASRS, (c, d) => c.SetComponent(d), registry);
                SpawnGroup(ps.robot_arms, _robotArmPool, ComponentType.RobotArm, (c, d) => c.SetComponent(d), registry);
                SpawnGroup(ps.processors, _processorPool, ComponentType.Processor, (c, d) => c.SetComponent(d), registry);
            }

            var infra = data.infrastructure;
            if (infra != null)
            {
                SpawnGroup(infra.queues, _queuePool, ComponentType.Queue, (c, d) => c.SetComponent(d), registry);
                SpawnGroup(infra.resources, _workerPool, ComponentType.Worker, (c, d) => c.SetComponent(d), registry);
            }

            return registry;
        }

        private void SpawnGroup<TComp, TData>(
            List<TData> items,
            GameObjectPool<TComp> pool,
            ComponentType type,
            Action<TComp, TData> setup,
            Dictionary<string, ComponentBase> registry)
            where TComp : ComponentBase
            where TData : ComponentDataBase
        {
            if (items == null) return;
            foreach (var item in items)
            {
                var comp = pool.GetItem(item.name);
                comp.componentType = type;
                setup(comp, item);
                _uiBinder.BindStandardComponent(comp);
                registry[item.name] = comp;
            }
        }

        public async UniTask<Dictionary<string, ComponentBase>> SpawnComponentsAsync(LogicDetailData data)
        {
            var registry = new Dictionary<string, ComponentBase>();

            int totalCount = CountAllItems(data);
            int spawnedCount = 0;
            int batchSize = 5;

            var ps = data.production_system;
            if (ps != null)
            {
                spawnedCount = await SpawnGroupAsync(ps.sources, _sourcePool, ComponentType.Source, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
                spawnedCount = await SpawnGroupAsync(ps.sinks, _sinkPool, ComponentType.Sink, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
                spawnedCount = await SpawnGroupAsync(ps.racks, _rackPool, ComponentType.Rack, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
                spawnedCount = await SpawnGroupAsync(ps.asrs, _asrsPool, ComponentType.ASRS, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
                spawnedCount = await SpawnGroupAsync(ps.robot_arms, _robotArmPool, ComponentType.RobotArm, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
                spawnedCount = await SpawnGroupAsync(ps.processors, _processorPool, ComponentType.Processor, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
            }

            var infra = data.infrastructure;
            if (infra != null)
            {
                spawnedCount = await SpawnGroupAsync(infra.queues, _queuePool, ComponentType.Queue, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
                spawnedCount = await SpawnGroupAsync(infra.resources, _workerPool, ComponentType.Worker, (c, d) => c.SetComponent(d), registry, spawnedCount, totalCount, batchSize);
            }

            return registry;
        }

        private async UniTask<int> SpawnGroupAsync<TComp, TData>(
            List<TData> items,
            GameObjectPool<TComp> pool,
            ComponentType type,
            Action<TComp, TData> setup,
            Dictionary<string, ComponentBase> registry,
            int spawnedCount,
            int totalCount,
            int batchSize)
            where TComp : ComponentBase
            where TData : ComponentDataBase
        {
            if (items == null) return spawnedCount;
            for (int i = 0; i < items.Count; i++)
            {
                var item = items[i];
                var comp = pool.GetItem(item.name);
                comp.componentType = type;
                setup(comp, item);
                _uiBinder.BindStandardComponent(comp);
                registry[item.name] = comp;

                spawnedCount++;
                if (totalCount > 0)
                    UILoadingBar.Percent = 0.7f * spawnedCount / totalCount;

                if (spawnedCount % batchSize == 0)
                    await UniTask.Yield();
            }
            return spawnedCount;
        }

        private int CountAllItems(LogicDetailData data)
        {
            int count = 0;
            var ps = data.production_system;
            if (ps != null)
            {
                if (ps.sources != null) count += ps.sources.Count;
                if (ps.sinks != null) count += ps.sinks.Count;
                if (ps.racks != null) count += ps.racks.Count;
                if (ps.asrs != null) count += ps.asrs.Count;
                if (ps.robot_arms != null) count += ps.robot_arms.Count;
                if (ps.processors != null) count += ps.processors.Count;
            }
            var infra = data.infrastructure;
            if (infra != null)
            {
                if (infra.queues != null) count += infra.queues.Count;
                if (infra.resources != null) count += infra.resources.Count;
            }
            return count;
        }

        public ComponentBase FindActiveItem(ComponentType type, string componentId)
        {
            switch (type)
            {
                case ComponentType.Source:    return _sourcePool?.FindActiveItem(componentId);
                case ComponentType.Sink:     return _sinkPool?.FindActiveItem(componentId);
                case ComponentType.Queue:    return _queuePool?.FindActiveItem(componentId);
                case ComponentType.Rack:     return _rackPool?.FindActiveItem(componentId);
                case ComponentType.ASRS:     return _asrsPool?.FindActiveItem(componentId);
                case ComponentType.RobotArm: return _robotArmPool?.FindActiveItem(componentId);
                case ComponentType.Processor:return _processorPool?.FindActiveItem(componentId);
                case ComponentType.Worker:   return _workerPool?.FindActiveItem(componentId);
                default: return null;
            }
        }
    }
}