Group party members within team allocators (#21)

NucleoidMC · Dec 8, 2024 · c84d27e · c84d27e
1 parent ac3fa4d
commit c84d27e
Showing 1 changed file with 21 additions and 3 deletions.
diff --git a/src/main/java/xyz/nucleoid/parties/PartyManager.java b/src/main/java/xyz/nucleoid/parties/PartyManager.java
@@ -15,6 +15,7 @@
 import java.util.HashSet;
 import java.util.UUID;
 import java.util.List;
+import java.util.stream.Collectors;
 
 public final class PartyManager {
     private static PartyManager instance;
@@ -35,11 +36,26 @@ public static void register() {
         GameEvents.COLLECT_PLAYERS_FOR_JOIN.register((gameSpace, player, additional) -> {
             var partyManager = PartyManager.get(player.server);
 
-            var members = partyManager.getPartyMembers(player);
+            var members = partyManager.getPartyMembers(player, true);
             Collections.shuffle((List<ServerPlayerEntity>) members);
 
             additional.addAll(members);
         });
+
+        GameEvents.TEAM_SELECTION_LOBBY_FINALIZE.register((gameSpace, allocator, players) -> {
+            var partyManager = PartyManager.get(gameSpace.getServer());
+
+            var ungroupedPlayers = players.stream().collect(Collectors.toCollection(HashSet::new));
+
+            for (ServerPlayerEntity player : players) {
+                if (ungroupedPlayers.contains(player)) {
+                    var members = partyManager.getPartyMembers(player, false);
+
+                    allocator.group(members);
+                    ungroupedPlayers.removeAll(members);
+                }
+            }
+        });
     }
 
     public static PartyManager get(MinecraftServer server) {
@@ -215,8 +231,10 @@ private Party createParty(PlayerRef owner) {
         return new Party(this.server, owner);
     }
 
-    public Collection<ServerPlayerEntity> getPartyMembers(ServerPlayerEntity player) {
-        var party = this.getOwnParty(PlayerRef.of(player));
+    public Collection<ServerPlayerEntity> getPartyMembers(ServerPlayerEntity player, boolean own) {
+        var ref = PlayerRef.of(player);
+        var party = own ? this.getOwnParty(ref) : this.getParty(ref);
+
         if (party != null) {
             return Lists.newArrayList(party.getMemberPlayers());
         } else {