Fixed CITItem not applying on tridents and spyglasses

Fixes #50

4 files changed, 68 insertions, 2 deletions

diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java b/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java
index b7d326b..0efd6f9 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java
@@ -6,6 +6,7 @@ import net.minecraft.entity.LivingEntity;
 import net.minecraft.item.ArmorItem;
 import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
+import net.minecraft.item.Items;
 import net.minecraft.util.Hand;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.registry.Registry;
@@ -77,11 +78,16 @@ public class ActiveCITs {
     public CITItem getCITItem(ItemStack stack, World world, LivingEntity entity) {
         Hand hand = entity != null && stack == entity.getOffHandStack() ? Hand.OFF_HAND : Hand.MAIN_HAND;
 
+        ((CITItem.Cached) (Object) stack).citresewn_setMojankCIT(false);
+
         List<CITItem> citItems = this.citItems.get(stack.getItem());
         if (citItems != null)
             for (CITItem citItem : citItems)
-                if (citItem.test(stack, hand, world, entity, true))
+                if (citItem.test(stack, hand, world, entity, true)) {
+                    if (stack.isOf(Items.TRIDENT) || stack.isOf(Items.SPYGLASS))
+                        ((CITItem.Cached) (Object) stack).citresewn_setMojankCIT(true);
                     return citItem;
+                }
         return null;
     }
 
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java
index c3aba64..ae097ca 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java
@@ -1,22 +1,39 @@
 package shcm.shsupercm.fabric.citresewn.mixin.cititem;
 
 import net.minecraft.client.MinecraftClient;
+import net.minecraft.client.render.VertexConsumerProvider;
+import net.minecraft.client.render.item.ItemModels;
 import net.minecraft.client.render.item.ItemRenderer;
 import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.json.ModelTransformation;
+import net.minecraft.client.util.ModelIdentifier;
+import net.minecraft.client.util.math.MatrixStack;
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.item.ItemStack;
+import net.minecraft.item.Items;
 import net.minecraft.world.World;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.ModifyVariable;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 import shcm.shsupercm.fabric.citresewn.CITResewn;
 import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
+import shcm.shsupercm.fabric.citresewn.pack.cits.CITItem;
+
+import java.lang.ref.WeakReference;
 
 @Mixin(ItemRenderer.class)
 public class ItemRendererMixin {
+    @Shadow @Final private ItemModels models;
+
+    private static WeakReference<BakedModel> mojankCITModel = null;
+
     @Inject(method = "getHeldItemModel", cancellable = true, at = @At("HEAD"))
-    public void getItemModel(ItemStack stack, World world, LivingEntity entity, int seed, CallbackInfoReturnable<BakedModel> cir) {
+    private void getItemModel(ItemStack stack, World world, LivingEntity entity, int seed, CallbackInfoReturnable<BakedModel> cir) {
         if (!CITResewnConfig.INSTANCE().enabled || CITResewn.INSTANCE.activeCITs == null)
             return;
 
@@ -24,4 +41,29 @@ public class ItemRendererMixin {
         if (citModel != null)
             cir.setReturnValue(citModel);
     }
+
+    @Inject(method = "renderItem(Lnet/minecraft/item/ItemStack;Lnet/minecraft/client/render/model/json/ModelTransformation$Mode;ZLnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;IILnet/minecraft/client/render/model/BakedModel;)V", at = @At("HEAD"))
+    private void fixMojankCITsContext(ItemStack stack, ModelTransformation.Mode renderMode, boolean leftHanded, MatrixStack matrices, VertexConsumerProvider vertexConsumers, int light, int overlay, BakedModel model, CallbackInfo ci) {
+        mojankCITModel = null;
+        if (((CITItem.Cached) (Object) stack).citresewn_isMojankCIT()) {
+            boolean bl = renderMode == ModelTransformation.Mode.GUI || renderMode == ModelTransformation.Mode.GROUND || renderMode == ModelTransformation.Mode.FIXED;
+            if (bl)
+                mojankCITModel = new WeakReference<>(model);
+            else { // rendered in hand model of trident/spyglass
+                if (stack.isOf(Items.TRIDENT))
+                    mojankCITModel = new WeakReference<>(this.models.getModelManager().getModel(new ModelIdentifier("minecraft:trident_in_hand#inventory")));
+                else if (stack.isOf(Items.SPYGLASS))
+                    mojankCITModel = new WeakReference<>(this.models.getModelManager().getModel(new ModelIdentifier("minecraft:spyglass_in_hand#inventory")));
+            }
+        } else
+            mojankCITModel = null;
+    }
+
+    @ModifyVariable(method = "renderItem(Lnet/minecraft/item/ItemStack;Lnet/minecraft/client/render/model/json/ModelTransformation$Mode;ZLnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;IILnet/minecraft/client/render/model/BakedModel;)V", at = @At(value = "LOAD", ordinal = 0, target = "Lnet/minecraft/client/render/model/BakedModel;getTransformation()Lnet/minecraft/client/render/model/json/ModelTransformation;"), argsOnly = true)
+    private BakedModel fixMojankCITs(BakedModel original) {
+        if (mojankCITModel != null)
+            return mojankCITModel.get();
+
+        return original;
+    }
 }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java
index f739fc8..4c888d8 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java
@@ -13,6 +13,11 @@ public class ItemStackMixin implements CITItem.Cached {
     private WeakReference<CITItem> citresewn_cachedCITItem = new WeakReference<>(null);
     private long citresewn_cacheTimeCITItem = 0;
 
+    /**
+     * Used by tridents and spy glasses to force their unique overrides to be applied correctly
+     */
+    private boolean citresewn_mojankCIT = false;
+
     @Override
     public CITItem citresewn_getCachedCITItem(Supplier<CITItem> realtime) {
         if (System.currentTimeMillis() - citresewn_cacheTimeCITItem >= CITResewnConfig.INSTANCE().cache_ms) {
@@ -22,4 +27,14 @@ public class ItemStackMixin implements CITItem.Cached {
 
         return citresewn_cachedCITItem.get();
     }
+
+    @Override
+    public boolean citresewn_isMojankCIT() {
+        return citresewn_mojankCIT;
+    }
+
+    @Override
+    public void citresewn_setMojankCIT(boolean mojankCIT) {
+        citresewn_mojankCIT = mojankCIT;
+    }
 }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java
index fc47096..217b043 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java
@@ -414,5 +414,8 @@ public class CITItem extends CIT {
 
     public interface Cached {
         CITItem citresewn_getCachedCITItem(Supplier<CITItem> realtime);
+
+        boolean citresewn_isMojankCIT();
+        void citresewn_setMojankCIT(boolean mojankCIT);
     }
 }