From 06571c1eb429c84440cb1f4c0536e4b070ec4fba Mon Sep 17 00:00:00 2001
From: Quantumlyy <ethoarx@gmail.com>
Date: Wed, 13 Dec 2023 22:41:09 +0100
Subject: [PATCH] check for queen in industrial apiary

---
 .../TileEntityApiamancersDrainerCommon.java   | 18 ++++++++---------
 .../TileEntityApiamancersDrainerGT.java       | 20 ++++++++++++++++---
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerCommon.java b/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerCommon.java
index d129d95..4704c5a 100644
--- a/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerCommon.java
+++ b/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerCommon.java
@@ -18,7 +18,6 @@ import forestry.apiculture.genetics.BeeGenome;
 import magicbees.bees.BeeManager;
 import magicbees.bees.BeeSpecies;
 import magicbees.main.Config;
-import magicbees.main.utils.LogHelper;
 import thaumcraft.api.aspects.Aspect;
 import thaumcraft.api.aspects.AspectList;
 import thaumcraft.api.aspects.IAspectContainer;
@@ -49,27 +48,24 @@ public class TileEntityApiamancersDrainerCommon extends TileEntity implements IA
                 TileEntity above = worldObj.getTileEntity(this.xCoord, this.yCoord + 1, this.zCoord);
                 IBeeHousing beeHousing = beeHousing(above);
                 if (beeHousing != null) {
-                    LogHelper.warn("is housing");
-
                     // This performs all the checks to see if the bee is a living queen and if the species conditions
                     // are met.
                     if (!canWork(beeHousing, above)) return;
-                    LogHelper.warn("can work");
 
-                    ItemStack queenStack = beeHousing.getBeeInventory().getQueen();
+                    // The beeRoot.isMember call will treat null the same as EnumBeeType.NONE which leads getSpecies to
+                    // return null.
+                    ItemStack queenStack = getQueen(beeHousing, above);
                     IAlleleBeeSpecies queenSpecies = BeeGenome.getSpecies(queenStack);
                     if (queenSpecies == null) return;
-                    LogHelper.warn("has species");
+
                     if (BeeManager.beeRoot.getType(queenStack) != EnumBeeType.QUEEN) return;
-                    LogHelper.warn("is queen");
 
                     if (Objects.equals(queenSpecies.getUID(), BeeSpecies.TC_ESSENTIA.getSpecies().getUID())) {
-                        LogHelper.warn("is essentia");
                         IBeeModifier modifier = BeeManager.beeRoot.createBeeHousingModifier(beeHousing);
                         IBee queen = BeeManager.beeRoot.getMember(queenStack);
                         IBeeGenome queenGenome = queen.getGenome();
-                        float productionMultiplier = modifier.getProductionModifier(queenGenome, 1.0F);
 
+                        float productionMultiplier = modifier.getProductionModifier(queenGenome, 1.0F);
                         int amount = (int) Math.ceil(Config.drainerAmount * Math.max(productionMultiplier, 1.0F));
 
                         addToContainer(aspect, amount);
@@ -91,6 +87,10 @@ public class TileEntityApiamancersDrainerCommon extends TileEntity implements IA
         return beekeepingLogic.canWork();
     }
 
+    protected ItemStack getQueen(IBeeHousing beeHousing, TileEntity te) {
+        return beeHousing.getBeeInventory().getQueen();
+    }
+
     private void drainQueen(IBeeHousing housing, IBeeModifier modifier, IBee queen) {
         float lifespanModifier = modifier.getLifespanModifier(queen.getGenome(), queen.getMate(), 1.0f);
         queen.age(housing.getWorld(), lifespanModifier);
diff --git a/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerGT.java b/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerGT.java
index 5827fed..15af8de 100644
--- a/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerGT.java
+++ b/src/main/java/magicbees/tileentity/TileEntityApiamancersDrainerGT.java
@@ -2,6 +2,7 @@ package magicbees.tileentity;
 
 import java.util.stream.Collectors;
 
+import net.minecraft.item.ItemStack;
 import net.minecraft.tileentity.TileEntity;
 import net.minecraft.util.EnumChatFormatting;
 
@@ -49,7 +50,6 @@ public class TileEntityApiamancersDrainerGT extends TileEntityApiamancersDrainer
 
         if (GTMetaTileEntity != null) {
             IMetaTileEntity underlyingMetaTileEntity = GTMetaTileEntity.getMetaTileEntity();
-
             if (!(underlyingMetaTileEntity instanceof GT_MetaTileEntity_IndustrialApiary)) return null;
 
             return (IBeeHousing) underlyingMetaTileEntity;
@@ -60,10 +60,24 @@ public class TileEntityApiamancersDrainerGT extends TileEntityApiamancersDrainer
 
     @Override
     protected boolean canWork(IBeeHousing beeHousing, TileEntity te) {
-        boolean canNormallyWork = super.canWork(beeHousing, te);
         BaseMetaTileEntity GTMetaTileEntity = GTTileEntity(te);
 
-        return GTMetaTileEntity != null ? GTMetaTileEntity.isActive() : canNormallyWork;
+        return GTMetaTileEntity != null ? GTMetaTileEntity.isActive() : super.canWork(beeHousing, te);
+    }
+
+    @Override
+    protected ItemStack getQueen(IBeeHousing beeHousing, TileEntity te) {
+        BaseMetaTileEntity GTMetaTileEntity = GTTileEntity(te);
+
+        if (GTMetaTileEntity != null) {
+            IMetaTileEntity underlyingMetaTileEntity = GTMetaTileEntity.getMetaTileEntity();
+            if (!(underlyingMetaTileEntity instanceof GT_MetaTileEntity_IndustrialApiary)) return null;
+            GT_MetaTileEntity_IndustrialApiary industrialApiary = (GT_MetaTileEntity_IndustrialApiary) underlyingMetaTileEntity;
+
+            return industrialApiary.getUsedQueen();
+        }
+
+        return super.getQueen(beeHousing, te);
     }
 
     private BaseMetaTileEntity GTTileEntity(TileEntity te) {