From 754960df2e69e3f938d2f9ab85d491946cf355b7 Mon Sep 17 00:00:00 2001
From: DrSmugleaf <DrSmugleaf@users.noreply.github.com>
Date: Thu, 13 Oct 2022 20:10:47 +0200
Subject: [PATCH] Fix actions bar and window not updating (#11899)

---
 .../Systems/Actions/ActionUIController.cs     | 133 ++++++++++++++----
 1 file changed, 108 insertions(+), 25 deletions(-)

diff --git a/Content.Client/UserInterface/Systems/Actions/ActionUIController.cs b/Content.Client/UserInterface/Systems/Actions/ActionUIController.cs
index b0fa714ce6..52b6c64b36 100644
--- a/Content.Client/UserInterface/Systems/Actions/ActionUIController.cs
+++ b/Content.Client/UserInterface/Systems/Actions/ActionUIController.cs
@@ -39,9 +39,9 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
     [Dependency] private readonly IEntityManager _entities = default!;
     [Dependency] private readonly IOverlayManager _overlays = default!;
 
-    [UISystemDependency] private readonly ActionsSystem _actionsSystem = default!;
-    [UISystemDependency] private readonly InteractionOutlineSystem _interactionOutline = default!;
-    [UISystemDependency] private readonly TargetOutlineSystem _targetOutline = default!;
+    [UISystemDependency] private readonly ActionsSystem? _actionsSystem = default;
+    [UISystemDependency] private readonly InteractionOutlineSystem? _interactionOutline = default;
+    [UISystemDependency] private readonly TargetOutlineSystem? _targetOutline = default;
 
     private const int DefaultPageIndex = 0;
     private ActionButtonContainer? _container;
@@ -100,8 +100,14 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
         _actionButton.OnPressed += ActionButtonPressed;
         _actionsBar.PageButtons.LeftArrow.OnPressed += OnLeftArrowPressed;
         _actionsBar.PageButtons.RightArrow.OnPressed += OnRightArrowPressed;
-        _actionsSystem.ActionReplaced += OnActionReplaced;
-        _actionsSystem.ActionsUpdated += OnActionsUpdated;
+
+        if (_actionsSystem != null)
+        {
+            _actionsSystem.ActionAdded += OnActionAdded;
+            _actionsSystem.ActionRemoved += OnActionRemoved;
+            _actionsSystem.ActionReplaced += OnActionReplaced;
+            _actionsSystem.ActionsUpdated += OnActionsUpdated;
+        }
 
         UpdateFilterLabel();
         SearchAndDisplay();
@@ -154,8 +160,13 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
 
     public void OnStateExited(GameplayState state)
     {
-        _actionsSystem.ActionReplaced -= OnActionReplaced;
-        _actionsSystem.ActionsUpdated -= OnActionsUpdated;
+        if (_actionsSystem != null)
+        {
+            _actionsSystem.ActionAdded -= OnActionAdded;
+            _actionsSystem.ActionRemoved -= OnActionRemoved;
+            _actionsSystem.ActionReplaced -= OnActionReplaced;
+            _actionsSystem.ActionsUpdated -= OnActionsUpdated;
+        }
 
         if (_window != null)
         {
@@ -189,7 +200,7 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
         if (CurrentPage[index] is not { } type)
             return;
 
-        _actionsSystem.TriggerAction(type);
+        _actionsSystem?.TriggerAction(type);
     }
 
     private void ChangePage(int index)
@@ -221,6 +232,78 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
         ChangePage(_currentPageIndex + 1);
     }
 
+    private void AppendAction(ActionType action)
+    {
+        if (_container == null)
+            return;
+
+        foreach (var button in _container.GetButtons())
+        {
+            if (button.Action != null)
+                continue;
+
+            SetAction(button, action);
+            return;
+        }
+
+        foreach (var page in _pages)
+        {
+            for (var i = 0; i < page.Size; i++)
+            {
+                var pageAction = page[i];
+                if (pageAction != null)
+                    continue;
+
+                page[i] = action;
+                return;
+            }
+        }
+    }
+
+    private void OnActionAdded(ActionType action)
+    {
+        foreach (var page in _pages)
+        {
+            for (var i = 0; i < page.Size; i++)
+            {
+                if (page[i] == action)
+                {
+                    return;
+                }
+            }
+        }
+
+        AppendAction(action);
+        SearchAndDisplay();
+    }
+
+    private void OnActionRemoved(ActionType action)
+    {
+        if (_container == null)
+            return;
+
+        foreach (var button in _container.GetButtons())
+        {
+            if (button.Action == action)
+            {
+                SetAction(button, null);
+            }
+        }
+
+        foreach (var page in _pages)
+        {
+            for (var i = 0; i < page.Size; i++)
+            {
+                if (page[i] == action)
+                {
+                    page[i] = null;
+                }
+            }
+        }
+
+        SearchAndDisplay();
+    }
+
     private void OnActionReplaced(ActionType existing, ActionType action)
     {
         if (_container == null)
@@ -285,8 +368,8 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
         return filter switch
         {
             Filters.Enabled => action.Enabled,
-            Filters.Item => action.Provider != null && action.Provider != _actionsSystem.PlayerActions?.Owner,
-            Filters.Innate => action.Provider == null || action.Provider == _actionsSystem.PlayerActions?.Owner,
+            Filters.Item => action.Provider != null && action.Provider != _actionsSystem?.PlayerActions?.Owner,
+            Filters.Innate => action.Provider == null || action.Provider == _actionsSystem?.PlayerActions?.Owner,
             Filters.Instant => action is InstantAction,
             Filters.Targeted => action is TargetedAction,
             _ => throw new ArgumentOutOfRangeException(nameof(filter), filter, null)
@@ -326,7 +409,7 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
         var search = _window.SearchBar.Text;
         var filters = _window.FilterButton.SelectedKeys;
 
-        IEnumerable<ActionType>? actions = _actionsSystem.PlayerActions?.Actions;
+        IEnumerable<ActionType>? actions = _actionsSystem?.PlayerActions?.Actions;
         actions ??= Array.Empty<ActionType>();
 
         if (filters.Count == 0 && string.IsNullOrWhiteSpace(search))
@@ -346,7 +429,7 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
             if (action.DisplayName.Contains(search, StringComparison.OrdinalIgnoreCase))
                 return true;
 
-            if (action.Provider == null || action.Provider == _actionsSystem.PlayerActions?.Owner)
+            if (action.Provider == null || action.Provider == _actionsSystem?.PlayerActions?.Owner)
                 return false;
 
             var name = _entities.GetComponent<MetaDataComponent>(action.Provider.Value).EntityName;
@@ -477,7 +560,7 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
                 return;
             }
 
-            _actionsSystem.TriggerAction(button.Action);
+            _actionsSystem?.TriggerAction(button.Action);
             _menuDragHelper.EndDrag();
         }
         else
@@ -559,18 +642,18 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
 
     public void OnSystemLoaded(ActionsSystem system)
     {
-        _actionsSystem.LinkActions += OnComponentLinked;
-        _actionsSystem.UnlinkActions += OnComponentUnlinked;
-        _actionsSystem.ClearAssignments += ClearActions;
-        _actionsSystem.AssignSlot += AssignSlots;
+        system.LinkActions += OnComponentLinked;
+        system.UnlinkActions += OnComponentUnlinked;
+        system.ClearAssignments += ClearActions;
+        system.AssignSlot += AssignSlots;
     }
 
     public void OnSystemUnloaded(ActionsSystem system)
     {
-        _actionsSystem.LinkActions -= OnComponentLinked;
-        _actionsSystem.UnlinkActions -= OnComponentUnlinked;
-        _actionsSystem.ClearAssignments -= ClearActions;
-        _actionsSystem.AssignSlot -= AssignSlots;
+        system.LinkActions -= OnComponentLinked;
+        system.UnlinkActions -= OnComponentUnlinked;
+        system.ClearAssignments -= ClearActions;
+        system.AssignSlot -= AssignSlots;
     }
 
     public override void FrameUpdate(FrameEventArgs args)
@@ -688,8 +771,8 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
 
         var range = entityAction.CheckCanAccess ? action.Range : -1;
 
-        _interactionOutline.SetEnabled(false);
-        _targetOutline.Enable(range, entityAction.CheckCanAccess, predicate, entityAction.Whitelist, null);
+        _interactionOutline?.SetEnabled(false);
+        _targetOutline?.Enable(range, entityAction.CheckCanAccess, predicate, entityAction.Whitelist, null);
     }
 
     /// <summary>
@@ -701,8 +784,8 @@ public sealed class ActionUIController : UIController, IOnStateChanged<GameplayS
             return;
 
         SelectingTargetFor = null;
-        _targetOutline.Disable();
-        _interactionOutline.SetEnabled(true);
+        _targetOutline?.Disable();
+        _interactionOutline?.SetEnabled(true);
 
         if (!_overlays.TryGetOverlay<ShowHandItemOverlay>(out var handOverlay) || handOverlay == null)
             return;