From 080937ae8d45c58d295c80b13a0b724740256295 Mon Sep 17 00:00:00 2001
From: CptMoore <39010654+CptMoore@users.noreply.github.com>
Date: Mon, 19 Jun 2023 21:09:33 +0200
Subject: [PATCH] Added support for armored components.

---
 .../CriticalEffects/CriticalEffectsCustom.cs  |  2 +
 source/Features/CriticalEffects/Criticals.cs  | 64 ++++++++++++++-----
 .../MechComponent_DamageComponent_Patch.cs    |  2 +-
 3 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/source/Features/CriticalEffects/CriticalEffectsCustom.cs b/source/Features/CriticalEffects/CriticalEffectsCustom.cs
index be45e297..6d64d317 100644
--- a/source/Features/CriticalEffects/CriticalEffectsCustom.cs
+++ b/source/Features/CriticalEffects/CriticalEffectsCustom.cs
@@ -25,6 +25,8 @@ public class CriticalEffectsCustom : SimpleCustomComponent, IAfterLoad, IIsDestr
     public string? CritFloatieMessage { get; set; } = null;
     public string? DestroyedFloatieMessage { get; set; } = null;
 
+    public bool IsArmored { get; set; } = false;
+
     protected virtual UnitType GetUnitType()
     {
         return UnitType.UNDEFINED;
diff --git a/source/Features/CriticalEffects/Criticals.cs b/source/Features/CriticalEffects/Criticals.cs
index 37904f1d..eaa904c5 100644
--- a/source/Features/CriticalEffects/Criticals.cs
+++ b/source/Features/CriticalEffects/Criticals.cs
@@ -16,15 +16,23 @@ internal class Criticals
     internal Criticals(MechComponent component)
     {
         this.component = component;
-        ce = new Lazy<CriticalEffectsCustom?>(FetchCriticalEffects);
+        if (actor == null)
+        {
+            throw new NullReferenceException($"{nameof(actor)} is null");
+        }
+
+        ce = new(FetchCriticalEffects);
     }
 
     private readonly MechComponent component;
-    private AbstractActor? actor => component.parent;
+    private AbstractActor actor => component.parent;
+
+    private bool IsLocationDestroyed => actor.StructureForLocation(component.Location) <= 0f;
 
     internal CriticalEffectsCustom? Effects => ce.Value;
     private readonly Lazy<CriticalEffectsCustom?> ce;
     private bool HasLinked => Effects?.LinkedStatisticName != null;
+    private bool IsArmored => Effects?.IsArmored ?? false;
     private CriticalEffectsCustom? FetchCriticalEffects()
     {
         var customs = component.componentDef.GetComponents<CriticalEffectsCustom>().ToList();
@@ -80,13 +88,8 @@ internal Criticals(MechComponent component)
         }
     }
 
-    internal void Hit(WeaponHitInfo hitInfo, ref ComponentDamageLevel damageLevel)
+    internal void Hit(WeaponHitInfo hitInfo, out ComponentDamageLevel damageLevel)
     {
-        if (actor == null)
-        {
-            return;
-        }
-
         SetHits(hitInfo, out damageLevel);
 
         if (damageLevel == ComponentDamageLevel.Destroyed)
@@ -97,18 +100,33 @@ internal void Hit(WeaponHitInfo hitInfo, ref ComponentDamageLevel damageLevel)
 
     private void SetHits(WeaponHitInfo hitInfo, out ComponentDamageLevel damageLevel)
     {
-        if (actor == null)
-        {
-            throw new NullReferenceException();
-        }
-
         int effectsMax, effectsPrev, effectsNext;
         {
             var compCritsMax = ComponentHitMax();
             var compCritsPrev = ComponentHitCount();
 
-            var locationDestroyed = actor.StructureForLocation(component.Location) <= 0f;
-            var possibleAddedHits = locationDestroyed ? compCritsMax : 1;
+            var possibleAddedHits = 1;
+            if (IsLocationDestroyed)
+            {
+                possibleAddedHits = compCritsMax;
+            }
+            else if (IsArmored)
+            {
+                var compCritsArmoredPrev = ComponentHitArmoredCount();
+                if (compCritsArmoredPrev < compCritsMax)
+                {
+                    var randomCache = actor.Combat.AttackDirector.GetRandomFromCache(hitInfo, 1);
+                    var isArmoredHit = compCritsArmoredPrev <= Mathf.RoundToInt(compCritsMax * randomCache[0]);
+                    if (isArmoredHit)
+                    {
+                        var compCritsArmoredNext = compCritsArmoredPrev + 1;
+                        ComponentHitArmoredCount(compCritsArmoredNext);
+                        damageLevel = component.DamageLevel;
+                        return;
+                    }
+                }
+            }
+
             var compCritsNext = Mathf.Min(compCritsMax, compCritsPrev + possibleAddedHits);
             var compCritsAdded = Mathf.Max(compCritsNext - compCritsPrev, 0);
 
@@ -186,6 +204,17 @@ private int ComponentHitCount(int? setHits = null)
         return stat.Get();
     }
 
+    private int ComponentHitArmoredCount(int? setHits = null)
+    {
+        var stat = component.StatCollection.MECriticalSlotsHitArmored();
+        stat.CreateIfMissing(); // move to Mech.init and remove "CreateIfMissing" from StatAdapter
+        if (setHits.HasValue)
+        {
+            stat.SetValue(setHits.Value);
+        }
+        return stat.Get();
+    }
+
     private int GroupHitCount(int? setHits = null)
     {
         if (actor == null)
@@ -474,4 +503,9 @@ internal static StatisticAdapter<int> MECriticalSlotsHit(this StatCollection sta
     {
         return new("MECriticalSlotsHit", statCollection, 0);
     }
+
+    internal static StatisticAdapter<int> MECriticalSlotsHitArmored(this StatCollection statCollection)
+    {
+        return new("MECriticalSlotsHitArmored", statCollection, 0);
+    }
 }
diff --git a/source/Features/CriticalEffects/Patches/MechComponent_DamageComponent_Patch.cs b/source/Features/CriticalEffects/Patches/MechComponent_DamageComponent_Patch.cs
index 87d4e83b..37a372f3 100644
--- a/source/Features/CriticalEffects/Patches/MechComponent_DamageComponent_Patch.cs
+++ b/source/Features/CriticalEffects/Patches/MechComponent_DamageComponent_Patch.cs
@@ -35,7 +35,7 @@ public static void Prefix(ref bool __runOriginal, MechComponent __instance, Weap
         }
 
         var mechComponent = __instance;
-        mechComponent.Criticals().Hit(hitInfo, ref damageLevel);
+        mechComponent.Criticals().Hit(hitInfo, out damageLevel);
     }
 
     [HarmonyPostfix]