Disease Stages But Epic

Content.IntegrationTests/Tests/DiseaseTest.cs

@@ -0,0 +1,48 @@
+using System.Threading.Tasks;
+using Content.Shared.Disease;
+using NUnit.Framework;
+using Robust.Shared.Prototypes;
+
+namespace Content.IntegrationTests.Tests;
+
+[TestFixture]
+public sealed class DiseaseTest
+{
+    /// <summary>
+    /// Asserts that a disease prototype has valid stages for its effects and cures.
+    /// </summary>
+    [Test]
+    public async Task Stages()
+    {
+        await using var pairTracker = await PoolManager.GetServerClient(new PoolSettings{NoClient = true});
+        var server = pairTracker.Pair.Server;
+
+        var protoManager = server.ResolveDependency<IPrototypeManager>();
+
+        await server.WaitAssertion(() =>
+        {
+            foreach (var proto in protoManager.EnumeratePrototypes<DiseasePrototype>())
+            {
+                var stages = proto.Stages;
+
+                foreach (var effect in proto.Effects)
+                {
+                    for (var i = 0; i < effect.Stages.Length; i++)
+                    {
+                        Assert.That(stages.Contains(i), $"Disease {proto.ID} has an effect with an incorrect stage, {i}!");
+                    }
+                }
+
+                foreach (var cure in proto.Cures)
+                {
+                    for (var i = 0; i < cure.Stages.Length; i++)
+                    {
+                        Assert.That(stages.Contains(i), $"Disease {proto.ID} has a cure with an incorrect stage, {i}!");
+                    }
+                }
+            }
+        });
+
+        await pairTracker.CleanReturnAsync();
+    }
+}

Content.Server/Disease/DiseaseSystem.cs

@@ -99,6 +99,7 @@ public override void Update(float frameTime)
                 {
                     var disease = carrierComp.Diseases[i];
                     disease.Accumulator += frameTime;
+                    disease.TotalAccumulator += frameTime;
 
                     if (disease.Accumulator < disease.TickTime) continue;
 
@@ -107,17 +108,29 @@ public override void Update(float frameTime)
                     var args = new DiseaseEffectArgs(carrierComp.Owner, disease, EntityManager);
                     disease.Accumulator -= disease.TickTime;
 
+                    int stage = 0; //defaults to stage 0 because you should always have one
+                    float lastThreshold = 0;
+                    for (var j = 0; j < disease.Stages.Count; j++)
+                    {
+                        if (disease.TotalAccumulator >= disease.Stages[j] &&
+                            disease.Stages[j] > lastThreshold)
+                        {
+                            lastThreshold = disease.Stages[j];
+                            stage = j;
+                        }
+                    }
+
                     foreach (var cure in disease.Cures)
                     {
-                        if (cure.Cure(args))
+                        if (cure.Stages.AsSpan().Contains(stage) && cure.Cure(args))
                             CureDisease(carrierComp, disease);
                     }
 
                     if (doEffects)
                     {
                         foreach (var effect in disease.Effects)
                         {
-                            if (_random.Prob(effect.Probability))
+                            if (effect.Stages.AsSpan().Contains(stage) && _random.Prob(effect.Probability))
                                 effect.Effect(args);
                         }
                     }

Content.Shared/Disease/DiseaseCure.cs

@@ -11,7 +11,12 @@ public abstract class DiseaseCure
         /// and false otherwise
         /// </summary>
         public abstract bool Cure(DiseaseEffectArgs args);
-
+        /// <summary>
+        /// What stages the cure applies to.
+        /// probably should be all, but go wild
+        /// </summary>
+        [DataField("stages")]
+        public readonly int[] Stages = { 0 };
         /// <summary>
         /// This is used by the disease diangoser machine
         /// to generate reports to tell people all of a disease's

Content.Shared/Disease/DiseaseEffect.cs

@@ -12,6 +12,11 @@ public abstract class DiseaseEffect
         [DataField("probability")]
         public float Probability = 1.0f;
         /// <summary>
+        ///     What stages this effect triggers on
+        /// </summary>
+        [DataField("stages")]
+        public readonly int[] Stages = { 0 };
+        /// <summary>
         /// What effect the disease will have.
         /// </summary>
         public abstract void Effect(DiseaseEffectArgs args);

Content.Shared/Disease/DiseasePrototype.cs

@@ -37,6 +37,19 @@ public sealed class DiseasePrototype : IPrototype, IInheritingPrototype
         /// </summary>
         public float Accumulator = 0f;
         /// <summary>
+        /// Since accumulator is reset with TickTime, this just tracks
+        /// the total amount of time a disease has been present.
+        /// </summary>
+        public float TotalAccumulator = 0f;
+        /// <summary>
+        /// Stores all the separate stages of the disease plus the time
+        /// thresholds for their activation
+        /// int: the disease stage (0 for baseline, 1, 2, etc.)
+        /// float: the time it takes for the stage to begin.
+        /// </summary>
+        [DataField("stages", serverOnly: true)]
+        public readonly List<float> Stages = new() { 0f };
+        /// <summary>
         /// List of effects the disease has that will
         /// run every second (by default anyway)
         /// </summary>