EntityComponentSystem/src/main/java/com/example/systems/SquadGenerationSystem.java

package com.example.systems;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Stream;

import com.example.components.Investment;
import com.example.components.Squad;
import com.example.ecs.Entity;
import com.example.ecs.Registry;

public class SquadGenerationSystem {
    private final Registry registry;

    public SquadGenerationSystem(Registry registry) {
        this.registry = registry;
    }

    public Entity generateSquad(Entity squadEntity, String name, int maxInvestment) {
        var investmentEntities = registry.getWithComponents(Investment.class);
        List<Entity> units = new ArrayList<>();
        for (var investmentEntity : investmentEntities) {
            var investment = registry.getComponent(investmentEntity, Investment.class);
            if (investment.forSquad().equals(squadEntity)) {
                units.addAll(createUnits(investment));
                registry.remove(investmentEntity);
            }
        }

        Squad squad = new Squad(name, units);
        registry.addComponent(squadEntity, squad);
        return squadEntity;
    }

    private List<Entity> createUnits(Investment investment) {
        int moneyRemaining = investment.amount();
        List<Entity> leader = new ArrayList<>();

        if (investment.witLeader()) {
            moneyRemaining -= investment.race().getLeaderCost();
            leader.add(investment.race().getFactory().createHero(registry));
        }

        var amount = moneyRemaining / investment.race().getCost();

        List<Entity> squad = Stream.generate(() -> 
            investment.race().getFactory().createUnit(registry))
        .limit(amount)
        .toList();
        var finalSquad = new ArrayList<Entity>();
        finalSquad.addAll(leader);
        finalSquad.addAll(squad);

        return finalSquad;
    }
}