Resolve ModTweaker #670: Centrifuge morb support

https://github.com/jaredlll08/ModTweaker/issues/670

src/main/java/com/blamejared/compat/thermalexpansion/Centrifuge.java

@@ -12,7 +12,13 @@ import net.minecraft.item.ItemStack;
 import net.minecraftforge.fluids.FluidStack;
 import stanhebben.zenscript.annotations.*;
 
+import cofh.core.init.CoreProps;
+import cofh.core.util.helpers.ItemHelper;
+import cofh.thermalexpansion.item.ItemMorb;
+import cofh.thermalfoundation.init.TFFluids;
+
 import java.util.Arrays;
+import java.util.ArrayList;
 
 @ZenClass("mods.thermalexpansion.Centrifuge")
 @ModOnly("thermalexpansion")
@@ -29,12 +35,26 @@ public class Centrifuge {
         }
         ModTweaker.LATE_ADDITIONS.add(new Add(InputHelper.toStacks(items), chances, InputHelper.toStack(input), energy, InputHelper.toFluid(fluid)));
     }
+
+    @ZenMethod
+    public static void addRecipeMob(String entityId, WeightedItemStack[] outputs, ILiquidStack fluid, int energy, int xp) {
+        IItemStack[] items = new IItemStack[outputs.length];
+        Integer[] chances = new Integer[outputs.length];
+        for(int i = 0; i < outputs.length; i++) {
+            items[i] = outputs[i].getStack();
+            chances[i] = (int) outputs[i].getPercent();
+        }
+        ModTweaker.LATE_ADDITIONS.add(new AddMob(InputHelper.toStacks(items), chances, entityId, energy, InputHelper.toFluid(fluid), xp));
+    }
     
     @ZenMethod
     public static void removeRecipe(IItemStack input) {
-        
         ModTweaker.LATE_REMOVALS.add(new Remove(InputHelper.toStack(input)));
-        
+    }
+
+    @ZenMethod
+    public static void removeRecipeMob(String entityId) {
+        ModTweaker.LATE_REMOVALS.add(new RemoveMob(entityId));
     }
     
     private static class Add extends BaseAction {
@@ -64,6 +84,59 @@ public class Centrifuge {
             return LogHelper.getStackDescription(input);
         }
     }
+
+    private static class AddMob extends BaseAction {
+        private ItemStack[] outputs;
+        private Integer[] chances;
+        private String entityId;
+        private Integer energy;
+        private FluidStack fluid;
+        private Integer xp;
+
+        public AddMob(ItemStack[] outputs, Integer[] chances, String entityId, Integer energy, FluidStack fluid, Integer xp) {
+            super("Centrifuge");
+            this.outputs = outputs;
+            this.chances = chances;
+            this.entityId = entityId;
+            this.energy = energy;
+            this.fluid = fluid;
+            this.xp = xp;
+        }
+
+        @Override
+        public void apply() {
+            if (!ItemMorb.validMobs.contains(entityId)) {
+                CraftTweakerAPI.logError("Not a valid morb entity: " + entityId);
+                return;
+            }
+
+            if (energy==null) { this.energy=CentrifugeManager.DEFAULT_ENERGY*2; }
+            if (fluid==null) { this.fluid=new FluidStack(TFFluids.fluidExperience, xp * CoreProps.MB_PER_XP);
+
+            // This is copied from CoFH CentrifugeManager.java:255 onwards
+            // It simplifies the need to create recipes for standard and
+            // reusable morbs. Modified slightly.
+            ArrayList<ItemStack> outputStandard = new ArrayList<ItemStack>(Arrays.asList(outputs));
+            ArrayList<ItemStack> outputReusable = new ArrayList<ItemStack>(Arrays.asList(outputs));
+
+            ArrayList<Integer> chanceStandard = new ArrayList<Integer>(Arrays.asList(chances));
+            ArrayList<Integer> chanceReusable = new ArrayList<Integer>(Arrays.asList(chances));
+
+            outputStandard.add(ItemHelper.cloneStack(ItemMorb.morbStandard));
+            outputReusable.add(ItemHelper.cloneStack(ItemMorb.morbReusable));
+
+            chanceStandard.add(ItemMorb.REUSE_CHANCE);
+            chanceReusable.add(100);
+
+            CentrifugeManager.addRecipeMob(energy, ItemMorb.setTag(ItemHelper.cloneStack(ItemMorb.morbStandard), this.entityId, false), outputStandard, chanceStandard, fluid);
+            CentrifugeManager.addRecipeMob(energy, ItemMorb.setTag(ItemHelper.cloneStack(ItemMorb.morbReusable), this.entityId, false), outputReusable, chanceReusable, fluid);
+        }
+
+        @Override
+        protected String getRecipeInfo() {
+            return entityId;
+        }
+    }
     
     private static class Remove extends BaseAction {
         
@@ -88,4 +161,39 @@ public class Centrifuge {
             return LogHelper.getStackDescription(input);
         }
     }
+
+    private static class RemoveMob extends BaseAction {
+        private String entityId;
+
+        public RemoveMob(String entityId) {
+            super("Centrifuge");
+            this.entityId = entityId;
+        }
+
+        @Override
+        public void apply() {
+            ItemStack standard = ItemMorb.setTag(ItemHelper.cloneStack(ItemMorb.morbStandard), entityId, false);
+            ItemStack reusable =  ItemMorb.setTag(ItemHelper.cloneStack(ItemMorb.morbReusable), entityId, false);
+
+            if (!CentrifugeManager.recipeExistsMob(standard) && !CentrifugeManager.recipeExistsMob(reusable)) {
+                CraftTweakerAPI.logError("No Centrifuge mob recipe exists for entity: " + entityId);
+                return;
+            }
+
+            // I think it's technically possible for there to be unmatched morbs
+            if (CentrifugeManager.recipeExistsMob(standard)) {
+                CentrifugeManager.removeRecipeMob(standard);
+            }
+
+            if (CentrifugeManager.recipeExistsMob(reusable)) {
+                CentrifugeManager.removeRecipeMob(reusable);
+            }
+        }
+
+        @Override
+        protected String getRecipeInfo() {
+            return entityId;
+        }
+    }
 }
+