From 1956f245095be717391c454240cba6bdacb61433 Mon Sep 17 00:00:00 2001
From: Logan Perkins <logan@git.lp-programming.com>
Date: Tue, 9 Jan 2024 02:47:48 -0800
Subject: [PATCH] Add ammoConsumedPerShot to ammosetdefs

---
 Source/CombatExtended/CombatExtended/Defs/AmmoSetDef.cs   | 2 ++
 .../StatWorkers/StatWorker_AmmoConsumedPerShotCount.cs    | 8 +++++---
 .../CombatExtended/CombatExtended/Verbs/Verb_ShootCE.cs   | 2 +-
 .../CombatExtended/Verbs/Verb_ShootMortarCE.cs            | 2 +-
 Source/VehiclesCompat/VehiclesCompat/VehiclesCompat.cs    | 1 +
 5 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/Source/CombatExtended/CombatExtended/Defs/AmmoSetDef.cs b/Source/CombatExtended/CombatExtended/Defs/AmmoSetDef.cs
index 31e66a2b42..f9e09b209c 100644
--- a/Source/CombatExtended/CombatExtended/Defs/AmmoSetDef.cs
+++ b/Source/CombatExtended/CombatExtended/Defs/AmmoSetDef.cs
@@ -16,6 +16,8 @@ public class AmmoSetDef : Def
 
         public AmmoSetDef similarTo;
 
+        public int ammoConsumedPerShot = 1;
+
         public override IEnumerable<StatDrawEntry> SpecialDisplayStats(StatRequest req)
         {
             foreach (StatDrawEntry entry in base.SpecialDisplayStats(req)) { yield return entry; }
diff --git a/Source/CombatExtended/CombatExtended/StatWorkers/StatWorker_AmmoConsumedPerShotCount.cs b/Source/CombatExtended/CombatExtended/StatWorkers/StatWorker_AmmoConsumedPerShotCount.cs
index 224805307b..0e9080355b 100644
--- a/Source/CombatExtended/CombatExtended/StatWorkers/StatWorker_AmmoConsumedPerShotCount.cs
+++ b/Source/CombatExtended/CombatExtended/StatWorkers/StatWorker_AmmoConsumedPerShotCount.cs
@@ -30,13 +30,15 @@ private Thing Gun(StatRequest req)
 
         public override bool ShouldShowFor(StatRequest req)
         {
-            return base.ShouldShowFor(req) && (GunDef(req)?.Verbs?
-                                               .Any(x => ((x as VerbPropertiesCE)?.ammoConsumedPerShotCount ?? 1) > 1) ?? false);
+            return base.ShouldShowFor(req) &&
+            (((GunDef(req)?.GetCompProperties<CompProperties_AmmoUser>() as CompProperties_AmmoUser)?.ammoSet.ammoConsumedPerShot != 1) ||
+             (GunDef(req)?.Verbs?.Any(x => ((x as VerbPropertiesCE)?.ammoConsumedPerShotCount ?? 1) > 1) ?? false));
         }
 
         public override float GetValueUnfinalized(StatRequest req, bool applyPostProcess = true)
         {
-            return ((VerbPropertiesCE)GunDef(req)?.Verbs?.FirstOrDefault(x => ((VerbPropertiesCE)x).ammoConsumedPerShotCount > 1))?.ammoConsumedPerShotCount ?? 1;
+            return ((GunDef(req)?.GetCompProperties<CompProperties_AmmoUser>() as CompProperties_AmmoUser)?.ammoSet?.ammoConsumedPerShot ?? 1 *
+            ((VerbPropertiesCE)GunDef(req)?.Verbs?.FirstOrDefault(x => ((VerbPropertiesCE)x).ammoConsumedPerShotCount > 1))?.ammoConsumedPerShotCount ?? 1);
         }
 
         public override string GetExplanationUnfinalized(StatRequest req, ToStringNumberSense numberSense)
diff --git a/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootCE.cs b/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootCE.cs
index 342a94b0d0..54bbb79d06 100644
--- a/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootCE.cs
+++ b/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootCE.cs
@@ -363,7 +363,7 @@ public override bool TryCastShot()
             //Reduce ammunition
             if (CompAmmo != null)
             {
-                if (!CompAmmo.TryReduceAmmoCount(VerbPropsCE.ammoConsumedPerShotCount))
+                if (!CompAmmo.TryReduceAmmoCount(CompAmmo.Props.ammoSet.ammoConsumedPerShot * VerbPropsCE.ammoConsumedPerShotCount))
                 {
                     return false;
                 }
diff --git a/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootMortarCE.cs b/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootMortarCE.cs
index 2f548cd217..354faef2e6 100644
--- a/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootMortarCE.cs
+++ b/Source/CombatExtended/CombatExtended/Verbs/Verb_ShootMortarCE.cs
@@ -227,7 +227,7 @@ public override bool TryCastShot()
             }
             if (CompAmmo != null)
             {
-                if (!CompAmmo.TryReduceAmmoCount(VerbPropsCE.ammoConsumedPerShotCount))
+                if (!CompAmmo.TryReduceAmmoCount(CompAmmo.Props.ammoSet.ammoConsumedPerShot * VerbPropsCE.ammoConsumedPerShotCount))
                 {
                     return false;
                 }
diff --git a/Source/VehiclesCompat/VehiclesCompat/VehiclesCompat.cs b/Source/VehiclesCompat/VehiclesCompat/VehiclesCompat.cs
index 57b9380810..25eeaad47f 100644
--- a/Source/VehiclesCompat/VehiclesCompat/VehiclesCompat.cs
+++ b/Source/VehiclesCompat/VehiclesCompat/VehiclesCompat.cs
@@ -105,6 +105,7 @@ public static IEnumerable<ThingDef> _GetUsedAmmo()
                                 cetddme._ammoSet = asd;
                                 var ammunition = vtd.ammunition = new ThingFilter();
                                 vtd.genericAmmo = false;
+                                vtd.chargePerAmmoCount = 1f / asd.ammoConsumedPerShot;
                                 HashSet<ThingDef> allowedAmmo = (HashSet<ThingDef>)ammunition.AllowedThingDefs;
 
                                 foreach (var al in asd.ammoTypes)