From 318ef70228a862491f9ea84c2633f95835584b59 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Fri, 5 May 2023 21:13:03 +0100
Subject: [PATCH] Fixed skeletons voiding picked up stacks

---
 .../unicopia/entity/ai/WantItTakeItGoal.java     | 16 +++++++++++-----
 .../enchantment/WantItNeedItEnchantment.java     |  8 ++++++++
 .../unicopia/mixin/MixinMobEntity.java           | 11 +++++++++++
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/entity/ai/WantItTakeItGoal.java b/src/main/java/com/minelittlepony/unicopia/entity/ai/WantItTakeItGoal.java
index a458ae21..2e92d5cb 100644
--- a/src/main/java/com/minelittlepony/unicopia/entity/ai/WantItTakeItGoal.java
+++ b/src/main/java/com/minelittlepony/unicopia/entity/ai/WantItTakeItGoal.java
@@ -8,7 +8,6 @@ import com.minelittlepony.unicopia.particle.ParticleUtils;
 import com.minelittlepony.unicopia.particle.UParticles;
 import net.minecraft.enchantment.EnchantmentHelper;
 import net.minecraft.entity.Entity;
-import net.minecraft.entity.Entity.RemovalReason;
 import net.minecraft.entity.EquipmentSlot;
 import net.minecraft.entity.ItemEntity;
 import net.minecraft.entity.LivingEntity;
@@ -50,6 +49,8 @@ public class WantItTakeItGoal extends BreakHeartGoal {
             speed *= 2;
         }
 
+        reach = Math.max(1.5, reach);
+
         mob.getNavigation().startMovingTo(target, speed);
 
         cooldown = Math.max(cooldown - 1, 0);
@@ -80,10 +81,15 @@ public class WantItTakeItGoal extends BreakHeartGoal {
                     ItemStack stack = item.getStack();
 
                     if (!item.isRemoved()) {
-                        mob.tryEquip(stack);
-                        mob.triggerItemPickedUpByEntityCriteria(item);
-                        mob.sendPickup(item, stack.getCount());
-                        item.remove(RemovalReason.DISCARDED);
+                        ItemStack collected = mob.tryEquip(stack.copy());
+                        if (!collected.isEmpty()) {
+                            mob.triggerItemPickedUpByEntityCriteria(item);
+                            mob.sendPickup(item, stack.getCount());
+                            stack.decrement(collected.getCount());
+                            if (stack.isEmpty()) {
+                                item.discard();
+                            }
+                        }
                     }
                 });
             }
diff --git a/src/main/java/com/minelittlepony/unicopia/item/enchantment/WantItNeedItEnchantment.java b/src/main/java/com/minelittlepony/unicopia/item/enchantment/WantItNeedItEnchantment.java
index d181a4d7..5624d39f 100644
--- a/src/main/java/com/minelittlepony/unicopia/item/enchantment/WantItNeedItEnchantment.java
+++ b/src/main/java/com/minelittlepony/unicopia/item/enchantment/WantItNeedItEnchantment.java
@@ -6,7 +6,9 @@ import com.minelittlepony.unicopia.particle.FollowingParticleEffect;
 import com.minelittlepony.unicopia.particle.ParticleUtils;
 import com.minelittlepony.unicopia.particle.UParticles;
 
+import net.minecraft.enchantment.EnchantmentHelper;
 import net.minecraft.entity.EquipmentSlot;
+import net.minecraft.item.ItemStack;
 
 public class WantItNeedItEnchantment extends SimpleEnchantment {
 
@@ -20,4 +22,10 @@ public class WantItNeedItEnchantment extends SimpleEnchantment {
             ParticleUtils.spawnParticles(new FollowingParticleEffect(UParticles.HEALTH_DRAIN, user.asEntity(), 0.2F), user.asEntity(), 1);
         }
     }
+
+    public static boolean prefersEquipment(ItemStack newStack, ItemStack oldStack) {
+        int newLevel = EnchantmentHelper.getLevel(UEnchantments.WANT_IT_NEED_IT, newStack);
+        int oldLevel = EnchantmentHelper.getLevel(UEnchantments.WANT_IT_NEED_IT, oldStack);
+        return newLevel > oldLevel;
+    }
 }
diff --git a/src/main/java/com/minelittlepony/unicopia/mixin/MixinMobEntity.java b/src/main/java/com/minelittlepony/unicopia/mixin/MixinMobEntity.java
index 4c18b1ec..6d5ef789 100644
--- a/src/main/java/com/minelittlepony/unicopia/mixin/MixinMobEntity.java
+++ b/src/main/java/com/minelittlepony/unicopia/mixin/MixinMobEntity.java
@@ -6,14 +6,17 @@ 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.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
 import com.minelittlepony.unicopia.entity.*;
 import com.minelittlepony.unicopia.entity.duck.RotatedView;
+import com.minelittlepony.unicopia.item.enchantment.WantItNeedItEnchantment;
 
 import net.minecraft.entity.EntityType;
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.ai.goal.GoalSelector;
 import net.minecraft.entity.mob.MobEntity;
+import net.minecraft.item.ItemStack;
 import net.minecraft.world.World;
 
 @Mixin(MobEntity.class)
@@ -41,4 +44,12 @@ abstract class MixinMobEntity extends LivingEntity implements Equine.Container<C
     public void afterTickAi(CallbackInfo into) {
         ((RotatedView)world).popRotation();
     }
+
+    @Inject(method = "prefersNewEquipment(Lnet/minecraft/item/ItemStack;Lnet/minecraft/item/ItemStack;)Z",
+            at = @At("HEAD"), cancellable = true)
+    private void onPrefersNewEquipment(ItemStack newStack, ItemStack oldStack, CallbackInfoReturnable<Boolean> info) {
+        if (WantItNeedItEnchantment.prefersEquipment(newStack, oldStack)) {
+            info.setReturnValue(true);
+        }
+    }
 }