riversong-code-showcase/Source/Riversong/Game/World/Agents/AgentManagerSystem.cs

using System.Collections.Generic;
using System.Runtime.CompilerServices;
using Cysharp.Threading.Tasks;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Pool;

namespace DanieleMarotta.RiversongCodeShowcase
{
    [Service(typeof(IAgentFactory))]
    [Service(typeof(IAgentVisualizationCollection))]
    [GameSystemGroup(typeof(DefaultAgentsSystemGroup))]
    public class AgentManagerSystem : GameSystem, IUpdatable, IAgentFactory, IAgentVisualizationCollection
    {
        [InjectService]
        private IEntityCollection _entityCollection;

        [InjectService]
        private IScene _scene;

        [InjectService]
        private ITileSpace _tileSpace;

        [InjectService]
        private IProductCatalog _productCatalog;

        [InjectService]
        private IProductStackFactory _productStackFactory;

        [InjectService]
        private World _world;

        private readonly Dictionary<int, (Agent, AgentVisualization)> _visualizations = new();

        private readonly HashSet<int> _pendingVisualizations = new();

        public AgentManagerSystem(IServiceLocator serviceLocator) : base(serviceLocator)
        {
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void InitializeAgentSource(IAgentSourceEntity source)
        {
            ref var agentSourceState = ref source.GetAgentSourceStateRW();
            agentSourceState.MaxAgentCount = agentSourceState.AgentIds.Capacity;
            agentSourceState.AgentCount.Length = (int)AgentJob.Count;
            agentSourceState.SpawnCooldown.Length = (int)AgentJob.Count;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool CanSpawnAgent(IAgentSourceEntity source, int maxAgentCount = int.MaxValue, AgentJob job = AgentJob.None)
        {
            ref var agentSourceState = ref source.GetAgentSourceStateRW();

            return agentSourceState.AgentIds.Length < agentSourceState.MaxAgentCount &&
                   agentSourceState.AgentCount[(int)job] < maxAgentCount &&
                   (job == AgentJob.None || agentSourceState.SpawnCooldown[(int)job] <= 0);
        }

        public Agent CreateVillager(AgentDefinition definition, IAgentSourceEntity source, float3 position, AgentJob job = AgentJob.None)
        {
            var villager = CreateAgent(definition, source, position, job);

            FinalizeAgentCreation(villager);

            return villager;
        }

        public Agent CreateCritter(CritterDefinition critterDefinition, IAgentSourceEntity source, float3 position)
        {
            var critter = CreateAgent(critterDefinition.AgentDefinition, source, position, AgentJob.None);

            ref var critterState = ref critter.GetCritterStateRW();
            critterState.IsCritter = true;
            critterState.CritterDefinitionId = critterDefinition.RuntimeId;

            FinalizeAgentCreation(critter);

            return critter;
        }

        private Agent CreateAgent(AgentDefinition definition, IAgentSourceEntity source, float3 position, AgentJob job)
        {
            var agent = _entityCollection.Create<Agent>();
            agent.Definition = definition;
            agent.HomeId = source.Id;
            agent.Position = position;
            agent.Heading = math.forward();
            agent.CarriedProductHandle = IProductCatalog.InvalidHandle;

            ref var path = ref agent.GetPathRW();
            path = TilePath.Initialize();

            ref var jobState = ref agent.GetJobStateRW();
            jobState.Job = job;

            ref var agentSourceState = ref source.GetAgentSourceStateRW();
            agentSourceState.AgentIds.Add(agent.Id);
            agentSourceState.AgentCount[(int)job]++;
            if (job != AgentJob.None) agentSourceState.SpawnCooldown[(int)job] = source.SpawnCooldown;

            return agent;
        }

        private void FinalizeAgentCreation(Agent agent)
        {
            _entityCollection.Add(agent);

            _ = CreateVisualizationAsync(agent.Definition, agent);
        }

        private async UniTask CreateVisualizationAsync(AgentDefinition definition, Agent agent)
        {
            var folder = _scene.SceneFolders.Agents;

            _pendingVisualizations.Add(agent.Id);
            var visualizationGameObject = await definition.Visualization.InstantiateAsync(agent.Position, Quaternion.identity, folder).ToUniTask();

            if (!_pendingVisualizations.Remove(agent.Id))
            {
                definition.Visualization.ReleaseInstance(visualizationGameObject);
                return;
            }

            var visualization = visualizationGameObject.GetComponent<AgentVisualization>();
            if (!visualization)
            {
                Debug.LogError($"Invalid agent visualization '{visualization.name}'");
                definition.Visualization.ReleaseInstance(visualizationGameObject);
                return;
            }

            _visualizations.Add(agent.Id, (agent, visualization));
        }

        public bool TryGetVisualization(int id, out AgentVisualization visualization)
        {
            if (_visualizations.TryGetValue(id, out var value))
            {
                (_, visualization) = value;
                return true;
            }
            visualization = null;
            return false;
        }

        public void Update()
        {
            DeSpawnAgents();

            UpdateVisualizations();
        }

        private void DeSpawnAgents()
        {
            var agents = _entityCollection.GetInternalEntityList(typeof(Agent));

            using var deSpawnScope = ListPool<Agent>.Get(out var deSpawn);
            foreach (Agent agent in agents)
                if (agent.LifecycleState == AgentLifecycleState.DeSpawning)
                    deSpawn.Add(agent);
            foreach (var agent in deSpawn) DeSpawnAgent(agent);

            foreach (Agent agent in agents)
                if (agent.IntentExecutionState == IntentExecutionState.EmptyQueue || !_entityCollection.Exists(agent.HomeId))
                    agent.LifecycleState = AgentLifecycleState.DeSpawning;
        }

        private void DeSpawnAgent(Agent agent)
        {
            var product = _productCatalog.GetProduct(agent.CarriedProductHandle);
            if (product)
            {
                var tile = _tileSpace.WorldToTile(agent.Position);

                ref var entityIdValue = ref _world.EntityIdMap.GetValueRW(tile);
                if (entityIdValue.BuildingId == agent.HomeId)
                {
                    var home = _entityCollection.Get<Building>(agent.HomeId);
                    ref var storagePolicy = ref home.GetProductStoragePolicyRW();
                    if (storagePolicy.IsTakingAllowed(agent.CarriedProductHandle))
                    {
                        ref var storage = ref home.GetStorageRW();
                        storage.Put(agent.CarriedProductHandle, 1);
                    }
                }
                else
                {
                    _productStackFactory.CreateOrMerge(tile, product, 1);
                }
            }

            agent.Dispose();

            _entityCollection.Remove(agent.Id);

            if (_pendingVisualizations.Remove(agent.Id)) return;

            var (_, visualization) = _visualizations[agent.Id];
            _visualizations.Remove(agent.Id);
            agent.Definition.Visualization.ReleaseInstance(visualization.gameObject);
        }

        private void UpdateVisualizations()
        {
            foreach (var (agent, visualization) in _visualizations.Values)
            {
                visualization.SetVisible(agent.LifecycleState == AgentLifecycleState.Live);

                if (agent.LifecycleState != AgentLifecycleState.Live) continue;

                var heading = math.normalizesafe(agent.Heading, math.forward());
                visualization.transform.SetPositionAndRotation(agent.Position, Quaternion.LookRotation(heading));
            }
        }
    }
}