Add enchanted staves to combat tab.

Provismet · Jan 2, 2024 · 45154e2 · 45154e2
1 parent 43baa52
commit 45154e2
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/main/java/com/provismet/AdditionalArmoury/registries/AAItemGroups.java b/src/main/java/com/provismet/AdditionalArmoury/registries/AAItemGroups.java
@@ -1,5 +1,6 @@
 package com.provismet.AdditionalArmoury.registries;
 
+import com.provismet.AdditionalArmoury.enchantments.staff.StaffEnchantment;
 import com.provismet.AdditionalArmoury.items.DaggerItem;
 
 import net.fabricmc.fabric.api.itemgroup.v1.ItemGroupEvents;
@@ -20,6 +21,19 @@ public static void register () {
         ItemGroupEvents.modifyEntriesEvent(ItemGroups.COMBAT).register(content -> content.addAfter(Items.NETHERITE_SWORD, AAItems.DAGGERS.toArray(new DaggerItem[0])));
         ItemGroupEvents.modifyEntriesEvent(ItemGroups.COMBAT).register(content -> content.addAfter(Items.TRIDENT, AAItems.STAFF));
 
+        ItemGroupEvents.modifyEntriesEvent(ItemGroups.COMBAT).register(content -> {
+            content.getContext().lookup().getOptionalWrapper(RegistryKeys.ENCHANTMENT).ifPresent(wrapper -> {
+                wrapper.streamEntries()
+                    .filter(entry -> entry.value() instanceof StaffEnchantment)
+                    .map(entry -> {
+                        ItemStack stack = AAItems.STAFF.getDefaultStack();
+                        stack.addEnchantment(entry.value(), 1);
+                        return stack;
+                    })
+                    .forEach(stack -> content.add(stack));
+            });
+        });
+
         ItemGroupEvents.modifyEntriesEvent(ItemGroups.COMBAT).register(content -> {
             content.getContext().lookup().getOptionalWrapper(RegistryKeys.POTION).ifPresent(wrapper -> {
                 AAItems.DAGGERS.forEach(dagger -> addPotions(content, wrapper, dagger, ItemGroup.StackVisibility.PARENT_AND_SEARCH_TABS));