ProjectDDD/Assets/_DDD/_Scripts/RestaurantState/FlowStates/RestaurantCustomerState.cs

using System.Collections.Generic;
using System.Threading.Tasks;
using Opsive.BehaviorDesigner.Runtime;
using Sirenix.OdinInspector;
using UnityEngine;
using UnityEngine.AddressableAssets;
using UnityEngine.ResourceManagement.AsyncOperations;

namespace DDD
{
    public class RestaurantCustomerState : ScriptableObject
    {
        private Dictionary<CustomerType, Subtree> _loadedSubtrees = new Dictionary<CustomerType, Subtree>();
        private Dictionary<CustomerType, AsyncOperationHandle<Subtree>> _subtreeHandles = new Dictionary<CustomerType, AsyncOperationHandle<Subtree>>();

        public async Task LoadCustomerBehaviorData()
        {
            var customerData = RestaurantData.Instance?.CustomerData;
            if (customerData?.CustomerBehaviorData == null)
            {
                Debug.LogError("[RestaurantCustomerState] RestaurantCustomerData or CustomerBehaviorData is null");
                return;
            }

            var loadTasks = new List<Task>();
            foreach (var behaviorPair in customerData.CustomerBehaviorData)
            {
                var customerType = behaviorPair.Key;
                var subtreeReference = behaviorPair.Value;
                
                if (_loadedSubtrees.ContainsKey(customerType))
                    continue; // Already loaded
                
                loadTasks.Add(LoadSubtreeAsync(customerType, subtreeReference));
            }

            await Task.WhenAll(loadTasks);
            Debug.Log($"[RestaurantCustomerState] Loaded {_loadedSubtrees.Count} customer behavior subtrees");
        }

        private async Task LoadSubtreeAsync(CustomerType customerType, AssetReference subtreeReference)
        {
            var handle = Addressables.LoadAssetAsync<Subtree>(subtreeReference);
            _subtreeHandles[customerType] = handle;
            
            await handle.Task;
            
            if (handle.Result != null)
            {
                _loadedSubtrees[customerType] = handle.Result;
                Debug.Log($"[RestaurantCustomerState] Loaded subtree for {customerType}");
            }
            else
            {
                Debug.LogError($"[RestaurantCustomerState] Failed to load subtree for {customerType}");
            }
        }

        public void UnloadCustomerBehaviorData()
        {
            foreach (var handle in _subtreeHandles.Values)
            {
                if (handle.IsValid())
                {
                    Addressables.Release(handle);
                }
            }
            
            _loadedSubtrees.Clear();
            _subtreeHandles.Clear();
            Debug.Log("[RestaurantCustomerState] Unloaded all customer behavior subtrees");
        }

        public Subtree GetLoadedSubtree(CustomerType customerType)
        {
            _loadedSubtrees.TryGetValue(customerType, out var subtree);
            return subtree;
        }

        public async Task<Subtree> GetOrLoadSubtree(CustomerType customerType)
        {
            if (IsSubtreeLoaded(customerType))
            {
                return GetLoadedSubtree(customerType);
            }
            else
            {
                var customerData = RestaurantData.Instance?.CustomerData;
                if (customerData?.CustomerBehaviorData == null || 
                    !customerData.CustomerBehaviorData.TryGetValue(customerType, out var subtreeReference))
                {
                    Debug.LogError($"[RestaurantCustomerState] No behavior data found for {customerType}");
                    return null;
                }

                await LoadSubtreeAsync(customerType, subtreeReference);
                return GetLoadedSubtree(customerType);
            }
        }

        public bool IsSubtreeLoaded(CustomerType customerType)
        {
            return _loadedSubtrees.ContainsKey(customerType);
        }
    }
}