From 217d081b29f7c15dccfa43d43d5f83121b75370b Mon Sep 17 00:00:00 2001
From: Nemanja <98561806+EmoGarbage404@users.noreply.github.com>
Date: Mon, 27 May 2024 18:43:17 -0400
Subject: [PATCH] Fix under-selecting antags (#28327)

Fix under selecting antags
---
 .../Antag/AntagSelectionSystem.API.cs         | 23 +++++++++++++++----
 Content.Server/Antag/AntagSelectionSystem.cs  |  2 +-
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/Content.Server/Antag/AntagSelectionSystem.API.cs b/Content.Server/Antag/AntagSelectionSystem.API.cs
index 3b5855cdf9..16d83c4f3f 100644
--- a/Content.Server/Antag/AntagSelectionSystem.API.cs
+++ b/Content.Server/Antag/AntagSelectionSystem.API.cs
@@ -52,12 +52,26 @@ public sealed partial class AntagSelectionSystem
     /// Gets the number of antagonists that should be present for a given rule based on the provided pool.
     /// A null pool will simply use the player count.
     /// </summary>
-    public int GetTargetAntagCount(Entity<AntagSelectionComponent> ent, AntagSelectionPlayerPool? pool = null)
+    public int GetTargetAntagCount(Entity<AntagSelectionComponent> ent, int? playerCount = null)
     {
         var count = 0;
         foreach (var def in ent.Comp.Definitions)
         {
-            count += GetTargetAntagCount(ent, pool, def);
+            count += GetTargetAntagCount(ent, playerCount, def);
+        }
+
+        return count;
+    }
+
+    public int GetTotalPlayerCount(IList<ICommonSession> pool)
+    {
+        var count = 0;
+        foreach (var session in pool)
+        {
+            if (session.Status is SessionStatus.Disconnected or SessionStatus.Zombie)
+                continue;
+
+            count++;
         }
 
         return count;
@@ -67,14 +81,13 @@ public sealed partial class AntagSelectionSystem
     /// Gets the number of antagonists that should be present for a given antag definition based on the provided pool.
     /// A null pool will simply use the player count.
     /// </summary>
-    public int GetTargetAntagCount(Entity<AntagSelectionComponent> ent, AntagSelectionPlayerPool? pool, AntagSelectionDefinition def)
+    public int GetTargetAntagCount(Entity<AntagSelectionComponent> ent, int? playerCount, AntagSelectionDefinition def)
     {
         // TODO ANTAG
         // make pool non-nullable
         // Review uses and ensure that people are INTENTIONALLY including players in the lobby if this is a mid-round
         // antag selection.
-        var poolSize = pool?.Count ?? _playerManager.Sessions
-            .Count(s => s.State.Status is not SessionStatus.Disconnected and not SessionStatus.Zombie);
+        var poolSize = playerCount ?? GetTotalPlayerCount(_playerManager.Sessions);
 
         // factor in other definitions' affect on the count.
         var countOffset = 0;
diff --git a/Content.Server/Antag/AntagSelectionSystem.cs b/Content.Server/Antag/AntagSelectionSystem.cs
index 5b6c891ddf..a86611bedb 100644
--- a/Content.Server/Antag/AntagSelectionSystem.cs
+++ b/Content.Server/Antag/AntagSelectionSystem.cs
@@ -203,7 +203,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
     public void ChooseAntags(Entity<AntagSelectionComponent> ent, IList<ICommonSession> pool, AntagSelectionDefinition def)
     {
         var playerPool = GetPlayerPool(ent, pool, def);
-        var count = GetTargetAntagCount(ent, playerPool, def);
+        var count = GetTargetAntagCount(ent, GetTotalPlayerCount(pool), def);
 
         for (var i = 0; i < count; i++)
         {
-- 
2.51.2