From babb279dd1154a9b719bb803efaaea08745794cd Mon Sep 17 00:00:00 2001
From: metalgearsloth <31366439+metalgearsloth@users.noreply.github.com>
Date: Mon, 4 Apr 2022 15:37:08 +1000
Subject: [PATCH] Disease artifact cleanup (#7359)

---
 Content.Server/Disease/DiseaseSystem.cs       |  8 +++-
 .../Components/DiseaseArtifactComponent.cs    | 17 ++------
 .../Effects/Systems/DiseaseArtifactSystem.cs  | 43 +++++++++++++------
 3 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/Content.Server/Disease/DiseaseSystem.cs b/Content.Server/Disease/DiseaseSystem.cs
index 17e8d27fcc..b4f553683b 100644
--- a/Content.Server/Disease/DiseaseSystem.cs
+++ b/Content.Server/Disease/DiseaseSystem.cs
@@ -332,9 +332,13 @@ namespace Content.Server.Disease
         /// rolls the dice to see if they get
         /// the disease.
         /// </summary>
-        public void TryInfect(DiseaseCarrierComponent carrier, DiseasePrototype? disease, float chance = 0.7f)
+        /// <param name="carrier">The target of the disease</param>
+        /// <param name="disease">The disease to apply</param>
+        /// <param name="chance">% chance of the disease being applied, before considering resistance</param>
+        /// <param name="forced">Bypass the disease's infectious trait.</param>
+        public void TryInfect(DiseaseCarrierComponent carrier, DiseasePrototype? disease, float chance = 0.7f, bool forced = false)
         {
-            if(disease is not { Infectious: true })
+            if(disease == null || !forced && !disease.Infectious)
                 return;
             var infectionChance = chance - carrier.DiseaseResist;
             if (infectionChance <= 0)
diff --git a/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Components/DiseaseArtifactComponent.cs b/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Components/DiseaseArtifactComponent.cs
index bcbe184030..16854f6eb6 100644
--- a/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Components/DiseaseArtifactComponent.cs
+++ b/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Components/DiseaseArtifactComponent.cs
@@ -8,14 +8,14 @@ namespace Content.Server.Xenoarchaeology.XenoArtifacts.Effects.Components;
 [RegisterComponent]
 public sealed class DiseaseArtifactComponent : Component
 {
-    public override string Name => "DiseaseArtifact";
     /// <summary>
     /// Disease the artifact will spawn
     /// If empty, picks a random one from its list
     /// </summary>
-    [DataField("disease", customTypeSerializer: typeof(PrototypeIdSerializer<DiseasePrototype>))]
+    [DataField("disease")]
     [ViewVariables(VVAccess.ReadWrite)]
-    public string SpawnDisease = string.Empty;
+    public DiseasePrototype? SpawnDisease;
+
     /// <summary>
     /// How far away it will check for people
     /// If empty, picks a random one from its list
@@ -23,15 +23,4 @@ public sealed class DiseaseArtifactComponent : Component
     [DataField("range")]
     [ViewVariables(VVAccess.ReadWrite)]
     public float Range = 5f;
-    [ViewVariables(VVAccess.ReadWrite)]
-    public DiseasePrototype ResolveDisease = default!;
-    [ViewVariables(VVAccess.ReadWrite)]
-    public readonly IReadOnlyList<string> ArtifactDiseases = new[]
-    {
-        "VanAusdallsRobovirus",
-        "OwOnavirus",
-        "BleedersBite",
-        "Ultragigacancer",
-        "AMIV"
-    };
 }
diff --git a/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Systems/DiseaseArtifactSystem.cs b/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Systems/DiseaseArtifactSystem.cs
index 520f953df4..407bf4ce95 100644
--- a/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Systems/DiseaseArtifactSystem.cs
+++ b/Content.Server/Xenoarchaeology/XenoArtifacts/Effects/Systems/DiseaseArtifactSystem.cs
@@ -6,6 +6,7 @@ using Content.Server.Disease.Components;
 using Robust.Shared.Random;
 using Robust.Shared.Prototypes;
 using Content.Shared.Interaction;
+using Robust.Shared.Map;
 
 namespace Content.Server.Xenoarchaeology.XenoArtifacts.Effects.Systems
 {
@@ -14,11 +15,23 @@ namespace Content.Server.Xenoarchaeology.XenoArtifacts.Effects.Systems
     /// </summary>
     public sealed class DiseaseArtifactSystem : EntitySystem
     {
-        [Dependency] private readonly EntityLookupSystem _lookup = default!;
         [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
         [Dependency] private readonly IRobustRandom _random = default!;
+        [Dependency] private readonly DiseaseSystem _disease = default!;
+        [Dependency] private readonly EntityLookupSystem _lookup = default!;
         [Dependency] private readonly SharedInteractionSystem _interactionSystem = default!;
 
+        // TODO: YAML Serializer won't catch this.
+        [ViewVariables(VVAccess.ReadWrite)]
+        public readonly IReadOnlyList<string> ArtifactDiseases = new[]
+        {
+            "VanAusdallsRobovirus",
+            "OwOnavirus",
+            "BleedersBite",
+            "Ultragigacancer",
+            "AMIV"
+        };
+
         public override void Initialize()
         {
             base.Initialize();
@@ -31,31 +44,37 @@ namespace Content.Server.Xenoarchaeology.XenoArtifacts.Effects.Systems
         /// </summary>
         private void OnMapInit(EntityUid uid, DiseaseArtifactComponent component, MapInitEvent args)
         {
-            if (component.SpawnDisease == string.Empty && component.ArtifactDiseases.Count != 0)
-            {
-                var diseaseName = _random.Pick(component.ArtifactDiseases);
+            if (component.SpawnDisease != null || ArtifactDiseases.Count == 0) return;
+            var diseaseName = _random.Pick(ArtifactDiseases);
 
-                component.SpawnDisease = diseaseName;
+            if (!_prototypeManager.TryIndex<DiseasePrototype>(diseaseName, out var disease))
+            {
+                Logger.ErrorS("disease", $"Invalid disease {diseaseName} selected from random diseases.");
+                return;
             }
 
-            if (_prototypeManager.TryIndex(component.SpawnDisease, out DiseasePrototype? disease) && disease != null)
-                component.ResolveDisease = disease;
+            component.SpawnDisease = disease;
         }
 
         /// <summary>
-        /// When activated, blasts everyone in LOS within 3 tiles
+        /// When activated, blasts everyone in LOS within n tiles
         /// with a high-probability disease infection attempt
         /// </summary>
         private void OnActivate(EntityUid uid, DiseaseArtifactComponent component, ArtifactActivatedEvent args)
         {
+            if (component.SpawnDisease == null) return;
+
             var xform = Transform(uid);
-            foreach (var entity in _lookup.GetEntitiesInRange(xform.MapID, xform.WorldPosition, 3f))
+            var carrierQuery = GetEntityQuery<DiseaseCarrierComponent>();
+
+            foreach (var entity in _lookup.GetEntitiesInRange(xform.Coordinates, component.Range))
             {
-                if (!_interactionSystem.InRangeUnobstructed(uid, entity, 3f))
+                if (!carrierQuery.TryGetComponent(entity, out var carrier)) continue;
+
+                if (!_interactionSystem.InRangeUnobstructed(uid, entity, component.Range))
                     continue;
 
-                if (TryComp<DiseaseCarrierComponent>(entity, out var carrier))
-                    EntitySystem.Get<DiseaseSystem>().TryInfect(carrier, component.ResolveDisease);
+                _disease.TryInfect(carrier, component.SpawnDisease, forced: true);
             }
         }
     }