Health alert tweaks

Content.Shared/Mobs/Systems/MobThresholdSystem.cs

@@ -5,7 +5,6 @@
 using Content.Shared.FixedPoint;
 using Content.Shared.Mobs.Components;
 using Robust.Shared.GameStates;
-using Robust.Shared.Utility;
 
 namespace Content.Shared.Mobs.Systems;
 
@@ -52,6 +51,38 @@ private void OnHandleState(EntityUid uid, MobThresholdsComponent component, ref
 
     #region Public API
 
+    /// <summary>
+    /// Gets the next available state for a mob.
+    /// </summary>
+    /// <param name="target">Target entity</param>
+    /// <param name="mobState">Supplied MobState</param>
+    /// <param name="nextState">The following MobState. Can be null if there isn't one.</param>
+    /// <param name="thresholdsComponent">Threshold Component Owned by the target</param>
+    /// <returns>True if the next mob state exists</returns>
+    public bool TryGetNextState(
+        EntityUid target,
+        MobState mobState,
+        [NotNullWhen(true)] out MobState? nextState,
+        MobThresholdsComponent? thresholdsComponent = null)
+    {
+        nextState = null;
+        if (!Resolve(target, ref thresholdsComponent))
+            return false;
+
+        MobState? min = null;
+        foreach (var state in thresholdsComponent.Thresholds.Values)
+        {
+            if (state <= mobState)
+                continue;
+
+            if (min == null || state < min)
+                min = state;
+        }
+
+        nextState = min;
+        return nextState != null;
+    }
+
     /// <summary>
     /// Get the Damage Threshold for the appropriate state if it exists
     /// </summary>
@@ -261,7 +292,7 @@ public void SetMobStateThreshold(EntityUid target, FixedPoint2 damage, MobState
             threshold.Thresholds.Remove(damageThreshold);
         }
         threshold.Thresholds[damage] = mobState;
-        Dirty(threshold);
+        Dirty(target, threshold);
         VerifyThresholds(target, threshold);
     }
 
@@ -291,7 +322,7 @@ public void SetAllowRevives(EntityUid uid, bool val, MobThresholdsComponent? com
         if (!Resolve(uid, ref component, false))
             return;
         component.AllowRevives = val;
-        Dirty(component);
+        Dirty(uid, component);
         VerifyThresholds(uid, component);
     }
 
@@ -344,42 +375,37 @@ private void UpdateAlerts(EntityUid target, MobState currentMobState, MobThresho
         if (!threshold.TriggersAlerts)
             return;
 
-        var dict = threshold.StateAlertDict;
-        var healthAlert = AlertType.HumanHealth;
-        var critAlert = AlertType.HumanCrit;
-        var deadAlert = AlertType.HumanDead;
+        if (!threshold.StateAlertDict.TryGetValue(currentMobState, out var currentAlert))
+        {
+            Log.Error($"No alert alert for mob state {currentMobState} for entity {ToPrettyString(target)}");
+            return;
+        }
 
-        dict.TryGetValue(MobState.Alive, out healthAlert);
-        dict.TryGetValue(MobState.Critical, out critAlert);
-        dict.TryGetValue(MobState.Dead, out deadAlert);
+        if (!_alerts.TryGet(currentAlert, out var alertPrototype))
+        {
+            Log.Error($"Invalid alert type {currentAlert}");
+            return;
+        }
 
-        switch (currentMobState)
+        if (alertPrototype.SupportsSeverity)
         {
-            case MobState.Alive:
+            var severity = _alerts.GetMinSeverity(currentAlert);
+            if (TryGetNextState(target, currentMobState, out var nextState, threshold) &&
+                TryGetPercentageForState(target, nextState.Value, damageable.TotalDamage, out var percentage))
             {
-                var severity = _alerts.GetMinSeverity(healthAlert);
-                if (TryGetIncapPercentage(target, damageable.TotalDamage, out var percentage))
-                {
-                    severity = (short) MathF.Floor(percentage.Value.Float() *
-                                                   _alerts.GetSeverityRange(healthAlert));
-                    severity += _alerts.GetMinSeverity(healthAlert);
-                }
-                _alerts.ShowAlert(target, healthAlert, severity);
-                break;
-            }
-            case MobState.Critical:
-            {
-                _alerts.ShowAlert(target, critAlert);
-                break;
-            }
-            case MobState.Dead:
-            {
-                _alerts.ShowAlert(target, deadAlert);
-                break;
+                percentage = FixedPoint2.Clamp(percentage.Value, 0, 1);
+
+                severity = (short) MathF.Round(
+                    MathHelper.Lerp(
+                        _alerts.GetMinSeverity(currentAlert),
+                        _alerts.GetMaxSeverity(currentAlert),
+                        percentage.Value.Float()));
             }
-            case MobState.Invalid:
-            default:
-                throw new ArgumentOutOfRangeException(nameof(currentMobState), currentMobState, null);
+            _alerts.ShowAlert(target, currentAlert, severity);
+        }
+        else
+        {
+            _alerts.ShowAlert(target, currentAlert);
         }
     }