ThingSetMaker_CountEnabledAmmoOnly

Patches/Core/ThingSetMakerDefs/ThingSetMakers_Reward.xml

@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<Patch>
+  <!--Make quest reward weapon include ammo-->
+  <Operation Class="PatchOperationAttributeSet">
+    <xpath>
+      Defs/ThingSetMakerDef[defName="Reward_ItemsStandard"]
+      /root/options/li[1]
+      /thingSetMaker/options/li[thingSetMaker[@Class="ThingSetMaker_MarketValue"]/fixedParams/filter/thingSetMakerTagsToAllow/li[text()="RewardStandardHighFreq" or text()="RewardStandardMidFreq" or text()="RewardStandardLowFreq" or text()="RewardStandardQualitySuper"]]
+      /thingSetMaker
+    </xpath>
+    <attribute>Class</attribute>
+    <value>CombatExtended.ThingSetMaker_MarketValueWithAmmo</value>
+  </Operation>
+
+  <Operation Class="PatchOperationAdd">
+    <xpath>
+      Defs/ThingSetMakerDef[defName="Reward_ItemsStandard"]
+      /root/options/li[1]
+      /thingSetMaker/options/li[thingSetMaker[@Class="CombatExtended.ThingSetMaker_MarketValueWithAmmo"]]
+      /thingSetMaker
+    </xpath>
+    <value>
+      <random>true</random>
+    </value>
+  </Operation>
+</Patch>

Source/CombatExtended/CombatExtended/CE_ThingSetMakerUtility.cs

@@ -0,0 +1,54 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Verse;
+
+namespace CombatExtended
+{
+    public static class CE_ThingSetMakerUtility
+    {
+        public static bool CanGenerate(ThingDef d, bool allowBasic, bool allowAdvanced)
+        {
+            return d is AmmoDef ammodef && (ammodef.tradeTags?.Contains(AmmoInjector.destroyWithAmmoDisabledTag) ?? false)
+                ? AmmoUtility.IsAmmoSystemActive(ammodef) && AdvancedOrBasic(ammodef, allowBasic, allowAdvanced)
+                : true;
+        }
+
+        public static bool AdvancedOrBasic(AmmoDef d, bool allowBasic, bool allowAdvanced)
+        {
+            return d.ammoClass == null || (d.ammoClass.advanced ? allowBasic : allowAdvanced);
+        }
+
+        public static ThingDef GetAmmoDef(CompProperties_AmmoUser comp, bool random, bool canGenerateAdvanced)
+        {
+            if (random)
+            {
+                var ammoDefs = comp.ammoSet.ammoTypes.Select(v => v.ammo).Where(d => canGenerateAdvanced || !d.ammoClass.advanced);
+                if (ammoDefs.EnumerableNullOrEmpty())
+                {
+                    ammoDefs = comp.ammoSet.ammoTypes.Select(v => v.ammo);
+                }
+                return ammoDefs.RandomElement();
+            }
+            return comp.ammoSet.ammoTypes.First().ammo;
+        }
+
+        public static void GenerateAmmoForWeapon(List<Thing> outThings, bool random, bool canGenerateAdvanced, IntRange magCount)
+        {
+            List<Thing> ammos = new List<Thing>();
+            foreach (var thing in outThings)
+            {
+                if (thing.TryGetComp<CompAmmoUser>() is CompAmmoUser ammoUser && ammoUser.UseAmmo)
+                {
+                    Thing ammo = ThingMaker.MakeThing(GetAmmoDef(ammoUser.Props, random, canGenerateAdvanced));
+                    ammo.stackCount = Math.Max(Math.Max(ammoUser.Props.AmmoGenPerMagOverride, ammoUser.Props.magazineSize), 1) * magCount.RandomInRange;
+                    ammos.Add(ammo);
+                }
+            }
+            foreach (var t in ammos)
+            {
+                outThings.Add(t);
+            }
+        }
+    }
+}

Source/CombatExtended/CombatExtended/ThingSetMaker_CountEnabledAmmoOnly.cs

@@ -0,0 +1,60 @@
+using System.Collections.Generic;
+using RimWorld;
+using Verse;
+
+namespace CombatExtended
+{
+    public class ThingSetMaker_CountEnabledAmmoOnly : ThingSetMaker_Count
+    {
+        bool basic = true;
+
+        bool advanced = false;
+
+        public override void Generate(ThingSetMakerParams parms, List<Thing> outThings)
+        {
+            parms.validator = d => CE_ThingSetMakerUtility.CanGenerate(d, basic, advanced);
+            base.Generate(parms, outThings);
+        }
+    }
+    public class ThingSetMaker_StackCountEnabledAmmoOnly : ThingSetMaker_StackCount
+    {
+        bool basic = true;
+
+        bool advanced = false;
+
+        public override void Generate(ThingSetMakerParams parms, List<Thing> outThings)
+        {
+            parms.validator = d => CE_ThingSetMakerUtility.CanGenerate(d, basic, advanced);
+            base.Generate(parms, outThings);
+        }
+    }
+    public class ThingSetMaker_CountWithAmmo : ThingSetMaker_Count
+    {
+        IntRange magCount = new IntRange(2, 5);
+
+        bool random = false;
+
+        bool canGenerateAdvanced = false;
+
+        public override void Generate(ThingSetMakerParams parms, List<Thing> outThings)
+        {
+            base.Generate(parms, outThings);
+            CE_ThingSetMakerUtility.GenerateAmmoForWeapon(outThings, random, canGenerateAdvanced, magCount);
+        }
+    }
+
+    public class ThingSetMaker_MarketValueWithAmmo : ThingSetMaker_MarketValue
+    {
+        IntRange magCount = new IntRange(2, 5);
+
+        bool random = false;
+
+        bool canGenerateAdvanced = false;
+
+        public override void Generate(ThingSetMakerParams parms, List<Thing> outThings)
+        {
+            base.Generate(parms, outThings);
+            CE_ThingSetMakerUtility.GenerateAmmoForWeapon(outThings, random, canGenerateAdvanced, magCount);
+        }
+    }
+}