From 17b014939a8f552d02cc102c9525820d22700303 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Tue, 24 Sep 2024 18:35:26 +0100
Subject: [PATCH] Polearms can no longer accept the sweeping enchant. Closes
 #408

---
 .../unicopia/item/PolearmItem.java            | 10 ++++++++-
 .../enchantment/CustomEnchantableItem.java    |  8 +++++++
 .../unicopia/mixin/MixinEnchantment.java      | 21 +++++++++++++++++++
 .../mixin/MixinEnchantmentHelper.java         |  9 ++++++++
 src/main/resources/unicopia.mixin.json        |  1 +
 5 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 src/main/java/com/minelittlepony/unicopia/item/enchantment/CustomEnchantableItem.java
 create mode 100644 src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantment.java

diff --git a/src/main/java/com/minelittlepony/unicopia/item/PolearmItem.java b/src/main/java/com/minelittlepony/unicopia/item/PolearmItem.java
index 5dc90908..85517dc5 100644
--- a/src/main/java/com/minelittlepony/unicopia/item/PolearmItem.java
+++ b/src/main/java/com/minelittlepony/unicopia/item/PolearmItem.java
@@ -7,14 +7,17 @@ import com.google.common.collect.Multimap;
 import com.minelittlepony.unicopia.UTags;
 import com.minelittlepony.unicopia.entity.Living;
 import com.minelittlepony.unicopia.entity.mob.UEntityAttributes;
+import com.minelittlepony.unicopia.item.enchantment.CustomEnchantableItem;
 
 import net.minecraft.block.*;
+import net.minecraft.enchantment.Enchantment;
+import net.minecraft.enchantment.Enchantments;
 import net.minecraft.entity.EquipmentSlot;
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.attribute.*;
 import net.minecraft.item.*;
 
-public class PolearmItem extends SwordItem {
+public class PolearmItem extends SwordItem implements CustomEnchantableItem {
     static final UUID ATTACK_RANGE_MODIFIER_ID = UUID.fromString("A7B3659C-AA74-469C-963A-09A391DCAA0F");
 
     private final Multimap<EntityAttribute, EntityAttributeModifier> attributeModifiers;
@@ -65,4 +68,9 @@ public class PolearmItem extends SwordItem {
 
         return true;
     }
+
+    @Override
+    public boolean isAcceptableEnchant(ItemStack stack, Enchantment enchantment) {
+        return enchantment != Enchantments.SWEEPING;
+    }
 }
diff --git a/src/main/java/com/minelittlepony/unicopia/item/enchantment/CustomEnchantableItem.java b/src/main/java/com/minelittlepony/unicopia/item/enchantment/CustomEnchantableItem.java
new file mode 100644
index 00000000..5fbec700
--- /dev/null
+++ b/src/main/java/com/minelittlepony/unicopia/item/enchantment/CustomEnchantableItem.java
@@ -0,0 +1,8 @@
+package com.minelittlepony.unicopia.item.enchantment;
+
+import net.minecraft.enchantment.Enchantment;
+import net.minecraft.item.ItemStack;
+
+public interface CustomEnchantableItem {
+    boolean isAcceptableEnchant(ItemStack stack, Enchantment enchantment);
+}
diff --git a/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantment.java b/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantment.java
new file mode 100644
index 00000000..5143711c
--- /dev/null
+++ b/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantment.java
@@ -0,0 +1,21 @@
+package com.minelittlepony.unicopia.mixin;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import com.minelittlepony.unicopia.item.enchantment.CustomEnchantableItem;
+
+import net.minecraft.enchantment.Enchantment;
+import net.minecraft.item.ItemStack;
+
+@Mixin(Enchantment.class)
+abstract class MixinEnchantment {
+    @Inject(method = "isAcceptableItem", at = @At("HEAD"), cancellable = true)
+    private void onIsAcceptableItem(ItemStack stack, CallbackInfoReturnable<Boolean> info) {
+        if (stack.getItem() instanceof CustomEnchantableItem item && !item.isAcceptableEnchant(stack, (Enchantment)(Object)this)) {
+            info.setReturnValue(false);
+        }
+    }
+}
diff --git a/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantmentHelper.java b/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantmentHelper.java
index 4458399c..0dd1aacf 100644
--- a/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantmentHelper.java
+++ b/src/main/java/com/minelittlepony/unicopia/mixin/MixinEnchantmentHelper.java
@@ -1,5 +1,7 @@
 package com.minelittlepony.unicopia.mixin;
 
+import java.util.List;
+
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
@@ -9,7 +11,9 @@ import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.enchantment.Enchantment;
 import net.minecraft.enchantment.EnchantmentHelper;
+import net.minecraft.enchantment.EnchantmentLevelEntry;
 import net.minecraft.entity.LivingEntity;
+import net.minecraft.item.ItemStack;
 
 @Mixin(EnchantmentHelper.class)
 abstract class MixinEnchantmentHelper {
@@ -23,4 +27,9 @@ abstract class MixinEnchantmentHelper {
             }
         });
     }
+
+    @Inject(method = "getPossibleEntries", at = @At("RETURN"))
+    private static void onGetPossibleEntries(int power, ItemStack stack, boolean treasureAllowed, CallbackInfoReturnable<List<EnchantmentLevelEntry>> info) {
+        info.getReturnValue().removeIf(entry -> !entry.enchantment.isAcceptableItem(stack));
+    }
 }
diff --git a/src/main/resources/unicopia.mixin.json b/src/main/resources/unicopia.mixin.json
index 17bd2121..be3b67d5 100644
--- a/src/main/resources/unicopia.mixin.json
+++ b/src/main/resources/unicopia.mixin.json
@@ -17,6 +17,7 @@
     "MixinChunkBlockLightProvider",
     "MutableBlockLightStorage",
     "MixinDamageSource",
+    "MixinEnchantment",
     "MixinEnchantmentHelper",
     "MixinFallLocation",
     "MixinEntity",