EE2Clone/Assets/Scripts/Systems/PopulationSystem.cs

using Unity.Burst;
using Unity.Collections;
using Unity.Entities;
using Unity.NetCode;
using EE2Clone.Components;
using EE2Clone.Core;

namespace EE2Clone.Systems
{
    /// <summary>
    /// Server-side: sums population-providing buildings per player and updates PopulationMax.
    /// </summary>
    [BurstCompile]
    [WorldSystemFilter(WorldSystemFilterFlags.ServerSimulation)]
    public partial struct PopulationSystem : ISystem
    {
        [BurstCompile]
        public void OnUpdate(ref SystemState state)
        {
            // Build a map of playerId → total pop capacity from buildings
            var popCapMap = new NativeHashMap<int, int>(GameConstants.MaxPlayers, Allocator.Temp);

            foreach (var (pop, owner) in
                     SystemAPI.Query<RefRO<ProvidesPopulation>, RefRO<OwnerPlayer>>()
                         .WithAll<BuildingTag>()
                         .WithNone<UnderConstructionTag>())
            {
                int playerId = owner.ValueRO.PlayerId;
                if (popCapMap.TryGetValue(playerId, out int current))
                    popCapMap[playerId] = current + pop.ValueRO.Amount;
                else
                    popCapMap[playerId] = GameConstants.StartingPopulationCap + pop.ValueRO.Amount;
            }

            // Count current population per player (units)
            var popCurrentMap = new NativeHashMap<int, int>(GameConstants.MaxPlayers, Allocator.Temp);

            foreach (var owner in SystemAPI.Query<RefRO<OwnerPlayer>>().WithAll<UnitTag>())
            {
                int playerId = owner.ValueRO.PlayerId;
                if (popCurrentMap.TryGetValue(playerId, out int current))
                    popCurrentMap[playerId] = current + 1;
                else
                    popCurrentMap[playerId] = 1;
            }

            // Update PlayerStateComponent
            foreach (var playerState in SystemAPI.Query<RefRW<PlayerStateComponent>>())
            {
                int pid = playerState.ValueRO.PlayerId;

                if (popCapMap.TryGetValue(pid, out int cap))
                    playerState.ValueRW.PopulationMax = System.Math.Min(cap, GameConstants.MaxPopulationCap);
                else
                    playerState.ValueRW.PopulationMax = GameConstants.StartingPopulationCap;

                if (popCurrentMap.TryGetValue(pid, out int current))
                    playerState.ValueRW.PopulationCurrent = current;
                else
                    playerState.ValueRW.PopulationCurrent = 0;
            }

            popCapMap.Dispose();
            popCurrentMap.Dispose();
        }
    }
}