Implemented enchantments and enchantmentLevels conditions

1 files changed, 19 insertions, 0 deletions

diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/cit/CITContext.java b/src/main/java/shcm/shsupercm/fabric/citresewn/cit/CITContext.java
index b2dd2cd..c5f1922 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/cit/CITContext.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/cit/CITContext.java
@@ -1,10 +1,18 @@
 package shcm.shsupercm.fabric.citresewn.cit;
 
 import net.minecraft.client.MinecraftClient;
+import net.minecraft.enchantment.EnchantmentHelper;
 import net.minecraft.entity.LivingEntity;
+import net.minecraft.item.EnchantedBookItem;
 import net.minecraft.item.ItemStack;
+import net.minecraft.item.Items;
+import net.minecraft.nbt.NbtCompound;
+import net.minecraft.nbt.NbtElement;
+import net.minecraft.util.Identifier;
 import net.minecraft.world.World;
 
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.Objects;
 
 public class CITContext {
@@ -12,12 +20,23 @@ public class CITContext {
     public final World world;
     public final LivingEntity entity;
 
+    private Map<Identifier, Integer> enchantments = null;
+
     public CITContext(ItemStack stack, World world, LivingEntity entity) {
         this.stack = stack;
         this.world = world == null ? MinecraftClient.getInstance().world : world;
         this.entity = entity;
     }
 
+    public Map<Identifier, Integer> enchantments() {
+        if (this.enchantments == null) {
+            this.enchantments = new LinkedHashMap<>();
+            for (NbtElement nbtElement : stack.isOf(Items.ENCHANTED_BOOK) ? EnchantedBookItem.getEnchantmentNbt(stack) : stack.getEnchantments())
+                this.enchantments.put(EnchantmentHelper.getIdFromNbt((NbtCompound) nbtElement), EnchantmentHelper.getLevelFromNbt((NbtCompound) nbtElement));
+        }
+        return this.enchantments;
+    }
+
     @Override
     public boolean equals(Object obj) {
         return obj instanceof CITContext that &&