From 4f8ea0c19b71c137655c32161de1245702e5e805 Mon Sep 17 00:00:00 2001
From: Leon Friedrich <60421075+ElectroJr@users.noreply.github.com>
Date: Mon, 19 Jun 2023 05:22:41 +1200
Subject: [PATCH] Fix damage visuals not resetting in some cases (#17399)

---
 Content.Client/Damage/DamageVisualsSystem.cs  | 54 ++++++++++---------
 .../Damage/Systems/DamageableSystem.cs        |  2 +-
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/Content.Client/Damage/DamageVisualsSystem.cs b/Content.Client/Damage/DamageVisualsSystem.cs
index 41f705c3f2..111ee2ddb5 100644
--- a/Content.Client/Damage/DamageVisualsSystem.cs
+++ b/Content.Client/Damage/DamageVisualsSystem.cs
@@ -351,22 +351,22 @@ public sealed class DamageVisualsSystem : VisualizerSystem<DamageVisualsComponen
         if (damageVisComp.Disabled)
             return;
 
-        HandleDamage(args.Component, damageVisComp);
+        HandleDamage(uid, args.Component, damageVisComp);
     }
 
-    private void HandleDamage(AppearanceComponent component, DamageVisualsComponent damageVisComp)
+    private void HandleDamage(EntityUid uid, AppearanceComponent component, DamageVisualsComponent damageVisComp)
     {
-        if (!TryComp(component.Owner, out SpriteComponent? spriteComponent)
-            || !TryComp(component.Owner, out DamageableComponent? damageComponent))
+        if (!TryComp(uid, out SpriteComponent? spriteComponent)
+            || !TryComp(uid, out DamageableComponent? damageComponent))
             return;
 
         if (damageVisComp.TargetLayers != null && damageVisComp.DamageOverlayGroups != null)
-            UpdateDisabledLayers(spriteComponent, component, damageVisComp);
+            UpdateDisabledLayers(uid, spriteComponent, component, damageVisComp);
 
         if (damageVisComp.Overlay && damageVisComp.DamageOverlayGroups != null && damageVisComp.TargetLayers == null)
             CheckOverlayOrdering(spriteComponent, damageVisComp);
 
-        if (AppearanceSystem.TryGetData<bool>(component.Owner, DamageVisualizerKeys.ForceUpdate, out var update, component)
+        if (AppearanceSystem.TryGetData<bool>(uid, DamageVisualizerKeys.ForceUpdate, out var update, component)
             && update)
         {
             ForceUpdateLayers(damageComponent, spriteComponent, damageVisComp);
@@ -376,11 +376,16 @@ public sealed class DamageVisualsSystem : VisualizerSystem<DamageVisualsComponen
         if (damageVisComp.TrackAllDamage)
         {
             UpdateDamageVisuals(damageComponent, spriteComponent, damageVisComp);
+            return;
         }
-        else if (AppearanceSystem.TryGetData<DamageVisualizerGroupData>(component.Owner, DamageVisualizerKeys.DamageUpdateGroups, out var data, component))
+
+        if (!AppearanceSystem.TryGetData<DamageVisualizerGroupData>(uid, DamageVisualizerKeys.DamageUpdateGroups,
+                out var data, component))
         {
-            UpdateDamageVisuals(data.GroupList, damageComponent, spriteComponent, damageVisComp);
+            data = new DamageVisualizerGroupData(Comp<DamageableComponent>(uid).DamagePerGroup.Keys.ToList());
         }
+
+        UpdateDamageVisuals(data.GroupList, damageComponent, spriteComponent, damageVisComp);
     }
 
     /// <summary>
@@ -389,29 +394,30 @@ public sealed class DamageVisualsSystem : VisualizerSystem<DamageVisualsComponen
     ///     layer will no longer be visible, or obtain
     ///     any damage updates.
     /// </summary>
-    private void UpdateDisabledLayers(SpriteComponent spriteComponent, AppearanceComponent component, DamageVisualsComponent damageVisComp)
+    private void UpdateDisabledLayers(EntityUid uid, SpriteComponent spriteComponent, AppearanceComponent component, DamageVisualsComponent damageVisComp)
     {
         foreach (var layer in damageVisComp.TargetLayerMapKeys)
         {
-            bool? layerStatus = null;
-            if (AppearanceSystem.TryGetData<bool>(component.Owner, layer, out var layerStateEnum, component))
-                layerStatus = layerStateEnum;
+            // I assume this gets set by something like body system if limbs are missing???
+            // TODO is this actually used by anything anywhere?
+            AppearanceSystem.TryGetData(uid, layer, out bool disabled, component);
 
-            if (layerStatus == null)
+            if (damageVisComp.DisabledLayers[layer] == disabled)
                 continue;
 
-            if (damageVisComp.DisabledLayers[layer] != (bool) layerStatus)
+            damageVisComp.DisabledLayers[layer] = disabled;
+            if (damageVisComp.TrackAllDamage)
             {
-                damageVisComp.DisabledLayers[layer] = (bool) layerStatus;
-                if (!damageVisComp.TrackAllDamage && damageVisComp.DamageOverlayGroups != null)
-                {
-                    foreach (var damageGroup in damageVisComp.DamageOverlayGroups!.Keys)
-                    {
-                        spriteComponent.LayerSetVisible($"{layer}{damageGroup}", damageVisComp.DisabledLayers[layer]);
-                    }
-                }
-                else if (damageVisComp.TrackAllDamage)
-                    spriteComponent.LayerSetVisible($"{layer}trackDamage", damageVisComp.DisabledLayers[layer]);
+                spriteComponent.LayerSetVisible($"{layer}trackDamage", !disabled);
+                continue;
+            }
+
+            if (damageVisComp.DamageOverlayGroups == null)
+                continue;
+
+            foreach (var damageGroup in damageVisComp.DamageOverlayGroups.Keys)
+            {
+                spriteComponent.LayerSetVisible($"{layer}{damageGroup}", !disabled);
             }
         }
     }
diff --git a/Content.Shared/Damage/Systems/DamageableSystem.cs b/Content.Shared/Damage/Systems/DamageableSystem.cs
index 98e0ce0117..b902fe292e 100644
--- a/Content.Shared/Damage/Systems/DamageableSystem.cs
+++ b/Content.Shared/Damage/Systems/DamageableSystem.cs
@@ -124,7 +124,7 @@ namespace Content.Shared.Damage
 
             if (EntityManager.TryGetComponent<AppearanceComponent>(uid, out var appearance) && damageDelta != null)
             {
-                var data = new DamageVisualizerGroupData(damageDelta.GetDamagePerGroup(_prototypeManager).Keys.ToList());
+                var data = new DamageVisualizerGroupData(component.DamagePerGroup.Keys.ToList());
                 _appearance.SetData(uid, DamageVisualizerKeys.DamageUpdateGroups, data, appearance);
             }
             RaiseLocalEvent(uid, new DamageChangedEvent(component, damageDelta, interruptsDoAfters, origin));