From b114f4a52313c91fe767bf1783efc9557da6fd95 Mon Sep 17 00:00:00 2001
From: 20kdc <asdd2808@gmail.com>
Date: Sun, 3 Oct 2021 12:33:28 +0100
Subject: [PATCH] Update UI when the AME node group is changed, fix AME
 connection issues (#4750)

* Update UI when the AME node group is changed (fixes core count indicator being wrong)

* Fix connection issues between AME controller and cores

* Fix #4365 by properly propagating and clearing AME core injection state

* Fixes #4364
---
 Content.Server/AME/AMENodeGroup.cs            | 34 +++++++++++++++----
 .../AME/Components/AMEControllerComponent.cs  | 11 ++++--
 .../AME/Components/AMEShieldComponent.cs      |  1 +
 .../Structures/Power/Generation/ame.yml       |  2 +-
 4 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/Content.Server/AME/AMENodeGroup.cs b/Content.Server/AME/AMENodeGroup.cs
index eb9acc62c3..c652095c3c 100644
--- a/Content.Server/AME/AMENodeGroup.cs
+++ b/Content.Server/AME/AMENodeGroup.cs
@@ -46,11 +46,6 @@ namespace Content.Server.AME
             foreach (var node in groupNodes)
             {
                 var nodeOwner = node.Owner;
-                if (nodeOwner.TryGetComponent(out AMEControllerComponent? controller))
-                {
-                    _masterController = controller;
-                }
-
                 if (nodeOwner.TryGetComponent(out AMEShieldComponent? shield))
                 {
                     var nodeNeighbors = grid.GetCellsInSquareArea(nodeOwner.Transform.Coordinates, 1)
@@ -61,6 +56,7 @@ namespace Content.Server.AME
                     {
                         _cores.Add(shield);
                         shield.SetCore();
+                        // Core visuals will be updated later.
                     }
                     else
                     {
@@ -68,10 +64,36 @@ namespace Content.Server.AME
                     }
                 }
             }
+
+            // Separate to ensure core count is correctly updated.
+            foreach (var node in groupNodes)
+            {
+                var nodeOwner = node.Owner;
+                if (nodeOwner.TryGetComponent(out AMEControllerComponent? controller))
+                {
+                    if (_masterController == null)
+                    {
+                        // Has to be the first one, as otherwise IsMasterController will return true on them all for this first update.
+                        _masterController = controller;
+                    }
+                    controller.OnAMENodeGroupUpdate();
+                }
+            }
+
+            UpdateCoreVisuals();
         }
 
-        public void UpdateCoreVisuals(int injectionAmount, bool injecting)
+        public void UpdateCoreVisuals()
         {
+            var injectionAmount = 0;
+            var injecting = false;
+
+            if (_masterController != null)
+            {
+                injectionAmount = _masterController.InjectionAmount;
+                injecting = _masterController.Injecting;
+            }
+
             var injectionStrength = CoreCount > 0 ? injectionAmount / CoreCount : 0;
 
             foreach (AMEShieldComponent core in _cores)
diff --git a/Content.Server/AME/Components/AMEControllerComponent.cs b/Content.Server/AME/Components/AMEControllerComponent.cs
index 1077eea4b8..6040fd00d1 100644
--- a/Content.Server/AME/Components/AMEControllerComponent.cs
+++ b/Content.Server/AME/Components/AMEControllerComponent.cs
@@ -27,7 +27,8 @@ namespace Content.Server.AME.Components
     public class AMEControllerComponent : SharedAMEControllerComponent, IActivate, IInteractUsing
     {
         [ViewVariables] private BoundUserInterface? UserInterface => Owner.GetUIOrNull(AMEControllerUiKey.Key);
-        [ViewVariables] private bool _injecting;
+        private bool _injecting;
+        [ViewVariables] public bool Injecting => _injecting;
         [ViewVariables] public int InjectionAmount;
 
         private AppearanceComponent? _appearance;
@@ -137,6 +138,12 @@ namespace Content.Server.AME.Components
             UpdateUserInterface();
         }
 
+        // Used to update core count
+        public void OnAMENodeGroupUpdate()
+        {
+            UpdateUserInterface();
+        }
+
         private AMEControllerBoundUserInterfaceState GetUserInterfaceState()
         {
             var jar = _jarSlot.ContainedEntity;
@@ -216,7 +223,7 @@ namespace Content.Server.AME.Components
                     break;
             }
 
-            GetAMENodeGroup()?.UpdateCoreVisuals(InjectionAmount, _injecting);
+            GetAMENodeGroup()?.UpdateCoreVisuals();
 
             UpdateUserInterface();
             ClickSound();
diff --git a/Content.Server/AME/Components/AMEShieldComponent.cs b/Content.Server/AME/Components/AMEShieldComponent.cs
index 4c64181a16..33fd8802d5 100644
--- a/Content.Server/AME/Components/AMEShieldComponent.cs
+++ b/Content.Server/AME/Components/AMEShieldComponent.cs
@@ -36,6 +36,7 @@ namespace Content.Server.AME.Components
         {
             _isCore = false;
             _appearance?.SetData(AMEShieldVisuals.Core, "isNotCore");
+            UpdateCoreVisuals(0, false);
         }
 
         public void UpdateCoreVisuals(int injectionStrength, bool injecting)
diff --git a/Resources/Prototypes/Entities/Structures/Power/Generation/ame.yml b/Resources/Prototypes/Entities/Structures/Power/Generation/ame.yml
index de36ed4ab0..e7096aa172 100644
--- a/Resources/Prototypes/Entities/Structures/Power/Generation/ame.yml
+++ b/Resources/Prototypes/Entities/Structures/Power/Generation/ame.yml
@@ -56,7 +56,7 @@
     examinable: true
     nodes:
       ame:
-        !type:CableDeviceNode
+        !type:AdjacentNode
         nodeGroupID: AMEngine
       input:
         !type:CableDeviceNode