From d21e4a1d4e6ad89e2baacd3e744c83278cc3b02c Mon Sep 17 00:00:00 2001
From: Slava0135 <40753025+Slava0135@users.noreply.github.com>
Date: Wed, 19 Jul 2023 00:11:43 +0300
Subject: [PATCH] Make power grid check event force APCs off (#17935)

* replace query with enumerator

* remove unnecessary imports

* prevent from toggling APC during event

* all entity query i guess
---
 .../PowerGridCheckDisabledComponent.cs        | 11 +++++++++
 .../Events/PowerGridCheckRule.cs              | 23 +++++++++++++++----
 2 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 Content.Server/StationEvents/Components/PowerGridCheckDisabledComponent.cs

diff --git a/Content.Server/StationEvents/Components/PowerGridCheckDisabledComponent.cs b/Content.Server/StationEvents/Components/PowerGridCheckDisabledComponent.cs
new file mode 100644
index 0000000000..90c6f2c22d
--- /dev/null
+++ b/Content.Server/StationEvents/Components/PowerGridCheckDisabledComponent.cs
@@ -0,0 +1,11 @@
+using Content.Server.StationEvents.Events;
+
+namespace Content.Server.StationEvents.Components;
+
+/// <summary>
+/// Added to APCs to prevent them from turning on until event end
+/// </summary>
+[RegisterComponent, Access(typeof(PowerGridCheckRule))]
+public sealed class PowerGridCheckDisabledComponent : Component
+{
+}
diff --git a/Content.Server/StationEvents/Events/PowerGridCheckRule.cs b/Content.Server/StationEvents/Events/PowerGridCheckRule.cs
index d24f04d800..acddc939c1 100644
--- a/Content.Server/StationEvents/Events/PowerGridCheckRule.cs
+++ b/Content.Server/StationEvents/Events/PowerGridCheckRule.cs
@@ -6,9 +6,7 @@ using Robust.Shared.Utility;
 using System.Threading;
 using Content.Server.Power.EntitySystems;
 using Timer = Robust.Shared.Timing.Timer;
-using System.Linq;
 using Content.Server.GameTicking.Rules.Components;
-using Robust.Shared.Random;
 using Content.Server.Station.Components;
 using Content.Server.StationEvents.Components;
 
@@ -19,6 +17,13 @@ namespace Content.Server.StationEvents.Events
     {
         [Dependency] private readonly ApcSystem _apcSystem = default!;
 
+        public override void Initialize()
+        {
+            base.Initialize();
+
+            SubscribeLocalEvent<PowerGridCheckDisabledComponent, ApcToggleMainBreakerAttemptEvent>(OnApcToggleMainBreaker);
+        }
+
         protected override void Started(EntityUid uid, PowerGridCheckRuleComponent component, GameRuleComponent gameRule, GameRuleStartedEvent args)
         {
             base.Started(uid, component, gameRule, args);
@@ -26,10 +31,11 @@ namespace Content.Server.StationEvents.Events
             if (!TryGetRandomStation(out var chosenStation))
                 return;
 
-            foreach (var (apc, transform) in EntityQuery<ApcComponent, TransformComponent>(true))
+            var query = AllEntityQuery<ApcComponent, TransformComponent>();
+            while (query.MoveNext(out var target, out var apc, out var transform))
             {
                 if (apc.MainBreakerEnabled && CompOrNull<StationMemberComponent>(transform.GridUid)?.Station == chosenStation)
-                    component.Powered.Add(apc.Owner);
+                    component.Powered.Add(target);
             }
 
             RobustRandom.Shuffle(component.Powered);
@@ -48,9 +54,10 @@ namespace Content.Server.StationEvents.Events
 
                 if (TryComp(entity, out ApcComponent? apcComponent))
                 {
-                    if(!apcComponent.MainBreakerEnabled)
+                    if (!apcComponent.MainBreakerEnabled)
                         _apcSystem.ApcToggleBreaker(entity, apcComponent);
                 }
+                RemComp<PowerGridCheckDisabledComponent>(entity);
             }
 
             // Can't use the default EndAudio
@@ -88,8 +95,14 @@ namespace Content.Server.StationEvents.Events
                     if (apcComponent.MainBreakerEnabled)
                         _apcSystem.ApcToggleBreaker(selected, apcComponent);
                 }
+                EnsureComp<PowerGridCheckDisabledComponent>(selected);
                 component.Unpowered.Add(selected);
             }
         }
+
+        private void OnApcToggleMainBreaker(EntityUid uid, PowerGridCheckDisabledComponent component, ref ApcToggleMainBreakerAttemptEvent args)
+        {
+            args.Cancelled = true;
+        }
     }
 }