finish augment api rewrite

src/main/java/com/portingdeadmods/modjam/api/augments/Augment.java

@@ -1,5 +1,7 @@
 package com.portingdeadmods.modjam.api.augments;
 
+import com.portingdeadmods.modjam.MJRegistries;
+import com.portingdeadmods.modjam.data.MJDataAttachments;
 import net.minecraft.core.HolderLookup;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.world.entity.player.Player;
@@ -43,6 +45,7 @@ public int getCooldown() {
 
     public void setCooldown(int cooldown) {
         this.cooldown = cooldown;
+        setChanged();
     }
 
     public void breakBlock(BlockEvent.BreakEvent event) {
@@ -71,13 +74,20 @@ public boolean isOnCooldown() {
         return getCooldown() > 0;
     }
 
+    // Call this, whenever NBT should be saved
+    protected final void setChanged() {
+        player.setData(MJDataAttachments.AUGMENT_DATA_CHANGED, MJRegistries.AUGMENT_SLOT.getId(augmentSlot));
+    }
+
     @Override
     public @UnknownNullability CompoundTag serializeNBT(HolderLookup.Provider provider) {
-        return new CompoundTag();
+        CompoundTag tag = new CompoundTag();
+        tag.putInt("cooldown", cooldown);
+        return tag;
     }
 
     @Override
     public void deserializeNBT(HolderLookup.Provider provider, CompoundTag nbt) {
-
+        this.cooldown = nbt.getInt("cooldown");
     }
 }

src/main/java/com/portingdeadmods/modjam/content/commands/GetAugmentCommand.java

@@ -6,12 +6,11 @@
 import com.portingdeadmods.modjam.api.augments.Augment;
 import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.api.augments.AugmentType;
-import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentSlotArgumentType;
 import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.commands.CommandSourceStack;
 import net.minecraft.commands.Commands;
 import net.minecraft.network.chat.Component;
-import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.entity.player.Player;
 
 // /modjam augments get <slot>

src/main/java/com/portingdeadmods/modjam/content/commands/GetAugmentCooldownCommand.java

@@ -0,0 +1,38 @@
+package com.portingdeadmods.modjam.content.commands;
+
+import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.context.CommandContext;
+import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.api.augments.AugmentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentSlotArgumentType;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
+import net.minecraft.commands.CommandSourceStack;
+import net.minecraft.commands.Commands;
+import net.minecraft.network.chat.Component;
+import net.minecraft.world.entity.player.Player;
+
+// /modjam augments get <slot>
+public class GetAugmentCooldownCommand {
+    public static void register(CommandDispatcher<CommandSourceStack> dispatcher) {
+        dispatcher.register(Commands.literal(ModJam.MODID)
+                .then(Commands.literal("augments")
+                        .then(Commands.literal("cooldown")
+                                .then(Commands.literal("get")
+                                        .then(Commands.argument("slot", AugmentSlotArgumentType.getInstance())
+                                                .executes(GetAugmentCooldownCommand::execute))))));
+    }
+
+    private static int execute(CommandContext<CommandSourceStack> ctx) {
+        Player player = ctx.getSource().getPlayer();
+        AugmentSlot augmentSlot = ctx.getArgument("slot", AugmentSlot.class);
+        Augment augmentBySlot = AugmentHelper.getAugmentBySlot(player, augmentSlot);
+        int cooldown = 0;
+        if (augmentBySlot != null) {
+            cooldown = augmentBySlot.getCooldown();
+            player.sendSystemMessage(Component.literal("Augment cooldown for slot '" + augmentSlot.getName() + "': " + cooldown));
+        }
+        return 1;
+    }
+}

src/main/java/com/portingdeadmods/modjam/content/commands/SetAugmentCommand.java

@@ -1,18 +1,17 @@
 package com.portingdeadmods.modjam.content.commands;
 
 import com.mojang.brigadier.CommandDispatcher;
-import com.mojang.brigadier.arguments.IntegerArgumentType;
 import com.mojang.brigadier.context.CommandContext;
 import com.portingdeadmods.modjam.ModJam;
 import com.portingdeadmods.modjam.api.augments.Augment;
 import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.api.augments.AugmentType;
-import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentSlotArgumentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentTypeArgumentType;
 import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.commands.CommandSourceStack;
 import net.minecraft.commands.Commands;
 import net.minecraft.network.chat.Component;
-import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.entity.player.Player;
 
 // /modjam augments set <slot> <augment>

src/main/java/com/portingdeadmods/modjam/content/commands/SetAugmentCooldownCommand.java

@@ -0,0 +1,44 @@
+package com.portingdeadmods.modjam.content.commands;
+
+import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.arguments.IntegerArgumentType;
+import com.mojang.brigadier.context.CommandContext;
+import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.api.augments.AugmentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentSlotArgumentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentTypeArgumentType;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
+import net.minecraft.commands.CommandSourceStack;
+import net.minecraft.commands.Commands;
+import net.minecraft.network.chat.Component;
+import net.minecraft.world.entity.player.Player;
+
+// /modjam augments set <slot> <augment>
+
+// TODO: Only set augments for slots that support them
+public class SetAugmentCooldownCommand {
+    public static void register(CommandDispatcher<CommandSourceStack> dispatcher) {
+        dispatcher.register(Commands.literal(ModJam.MODID)
+                .then(Commands.literal("augments")
+                        .then(Commands.literal("cooldown")
+                                .then(Commands.literal("set")
+                                        .then(Commands.argument("slot", AugmentSlotArgumentType.getInstance())
+                                                .then(Commands.argument("cooldown", IntegerArgumentType.integer())
+                                                        .executes(SetAugmentCooldownCommand::execute)))))));
+    }
+
+    private static int execute(CommandContext<CommandSourceStack> ctx) {
+        Player player = ctx.getSource().getPlayer();
+        AugmentSlot slot = ctx.getArgument("slot", AugmentSlot.class);
+        int cooldown = ctx.getArgument("cooldown", int.class);
+        Augment augmentBySlot = AugmentHelper.getAugmentBySlot(player, slot);
+        if (augmentBySlot != null) {
+            augmentBySlot.setCooldown(cooldown);
+            player.sendSystemMessage(Component.literal("Set augment cooldown in slot '" + slot.getName() + "' to: " + cooldown));
+        }
+        return 1;
+    }
+
+}

src/main/java/com/portingdeadmods/modjam/content/commands/AugmentSlotArgumentType.java → src/main/java/com/portingdeadmods/modjam/content/commands/arguments/AugmentSlotArgumentType.java

@@ -1,4 +1,4 @@
-package com.portingdeadmods.modjam.content.commands;
+package com.portingdeadmods.modjam.content.commands.arguments;
 
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.arguments.ArgumentType;
@@ -10,7 +10,6 @@
 import com.portingdeadmods.modjam.MJRegistries;
 import com.portingdeadmods.modjam.ModJam;
 import com.portingdeadmods.modjam.api.augments.AugmentSlot;
-import com.portingdeadmods.modjam.api.augments.AugmentType;
 import net.minecraft.commands.SharedSuggestionProvider;
 import net.minecraft.network.chat.Component;
 import net.minecraft.resources.ResourceKey;

src/main/java/com/portingdeadmods/modjam/content/commands/AugmentTypeArgumentType.java → src/main/java/com/portingdeadmods/modjam/content/commands/arguments/AugmentTypeArgumentType.java

@@ -1,15 +1,13 @@
-package com.portingdeadmods.modjam.content.commands;
+package com.portingdeadmods.modjam.content.commands.arguments;
 
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.arguments.ArgumentType;
 import com.mojang.brigadier.context.CommandContext;
-import com.mojang.brigadier.exceptions.CommandExceptionType;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
 import com.mojang.brigadier.exceptions.DynamicCommandExceptionType;
 import com.mojang.brigadier.suggestion.Suggestions;
 import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import com.portingdeadmods.modjam.MJRegistries;
-import com.portingdeadmods.modjam.ModJam;
 import com.portingdeadmods.modjam.api.augments.AugmentType;
 import net.minecraft.commands.SharedSuggestionProvider;
 import net.minecraft.network.chat.Component;

src/main/java/com/portingdeadmods/modjam/data/MJDataAttachments.java

@@ -25,4 +25,7 @@ public final class MJDataAttachments {
             "augments_extra_data", () -> AttachmentType.<Map<AugmentSlot, CompoundTag>>builder(Collections::emptyMap)
                     .serialize(AugmentCodecs.AUGMENTS_EXTRA_DATA_CODEC).copyOnDeath().build()
     );
+    public static final Supplier<AttachmentType<Integer>> AUGMENT_DATA_CHANGED = ATTACHMENTS.register(
+            "augment_data_changed", () -> AttachmentType.builder(() -> -1).build()
+    );
 }

src/main/java/com/portingdeadmods/modjam/events/MJEvents.java

@@ -2,28 +2,36 @@
 
 import com.portingdeadmods.modjam.MJRegistries;
 import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.Augment;
 import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.api.augments.AugmentType;
-import com.portingdeadmods.modjam.content.commands.AugmentSlotArgumentType;
-import com.portingdeadmods.modjam.content.commands.AugmentTypeArgumentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentSlotArgumentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentTypeArgumentType;
 import com.portingdeadmods.modjam.content.recipes.ItemEtchingRecipe;
+import com.portingdeadmods.modjam.data.MJDataAttachments;
 import com.portingdeadmods.modjam.registries.MJFluidTypes;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import com.portingdeadmods.modjam.utils.ParticlesUtils;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
 import net.minecraft.core.Registry;
 import net.minecraft.core.particles.ParticleTypes;
+import net.minecraft.nbt.CompoundTag;
 import net.minecraft.world.entity.item.ItemEntity;
+import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.crafting.RecipeHolder;
 import net.minecraft.world.item.crafting.SingleRecipeInput;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.phys.Vec3;
 import net.neoforged.bus.api.SubscribeEvent;
 import net.neoforged.fml.common.EventBusSubscriber;
+import net.neoforged.neoforge.event.entity.player.PlayerEvent;
 import net.neoforged.neoforge.event.tick.EntityTickEvent;
+import net.neoforged.neoforge.event.tick.PlayerTickEvent;
 import net.neoforged.neoforge.registries.RegisterEvent;
 
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -45,6 +53,35 @@ public static void onItemEntityTick(EntityTickEvent.Post event) {
             }
         }
 
+        @SubscribeEvent
+        public static void onPlayerLoggedIn(PlayerEvent.PlayerLoggedInEvent event) {
+            Player player = event.getEntity();
+            Map<AugmentSlot, Augment> augments = player.getData(MJDataAttachments.AUGMENTS);
+            Map<AugmentSlot, CompoundTag> augmentsExtraData = player.getData(MJDataAttachments.AUGMENTS_EXTRA_DATA);
+            for (AugmentSlot augmentSlot : augments.keySet()) {
+                Augment augment = augments.get(augmentSlot);
+                augment.setPlayer(player);
+                CompoundTag nbt = augmentsExtraData.get(augmentSlot);
+                if (nbt != null) {
+                    augment.deserializeNBT(player.level().registryAccess(), nbt);
+                }
+            }
+        }
+
+        @SubscribeEvent
+        public static void onPlayerTick(PlayerTickEvent.Post event) {
+            Player player = event.getEntity();
+            int changedIndex = player.getData(MJDataAttachments.AUGMENT_DATA_CHANGED);
+            if (changedIndex != -1) {
+                Map<AugmentSlot, Augment> augments = AugmentHelper.getAugments(player);
+                Map<AugmentSlot, CompoundTag> augmentsExtraData = AugmentHelper.getAugmentsData(player);
+                AugmentSlot changedSlot = MJRegistries.AUGMENT_SLOT.byId(changedIndex);
+                CompoundTag tag = augments.get(changedSlot).serializeNBT(player.level().registryAccess());
+                AugmentHelper.setAugmentExtraData(player, changedSlot, tag);
+                player.setData(MJDataAttachments.AUGMENT_DATA_CHANGED, -1);
+            }
+        }
+
         private static void processItemEtching(ItemEntity itemEntity, Level level) {
             ItemStack stack = itemEntity.getItem();
 

src/main/java/com/portingdeadmods/modjam/registries/MJArgumentTypes.java

@@ -1,15 +1,13 @@
 package com.portingdeadmods.modjam.registries;
 
 import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.content.commands.AugmentSlotArgumentType;
-import com.portingdeadmods.modjam.content.commands.AugmentTypeArgumentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentSlotArgumentType;
+import com.portingdeadmods.modjam.content.commands.arguments.AugmentTypeArgumentType;
 import net.minecraft.commands.synchronization.ArgumentTypeInfo;
 import net.minecraft.commands.synchronization.ArgumentTypeInfos;
 import net.minecraft.commands.synchronization.SingletonArgumentInfo;
 import net.minecraft.core.registries.Registries;
 import net.neoforged.neoforge.registries.DeferredRegister;
-import top.theillusivec4.curios.api.CuriosApi;
-import top.theillusivec4.curios.server.command.CurioArgumentType;
 
 import java.util.function.Supplier;
 

src/main/java/com/portingdeadmods/modjam/registries/MJCommands.java

@@ -1,8 +1,12 @@
 package com.portingdeadmods.modjam.registries;
 
+import com.mojang.brigadier.CommandDispatcher;
 import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.content.commands.GetAugmentCooldownCommand;
 import com.portingdeadmods.modjam.content.commands.SetAugmentCommand;
 import com.portingdeadmods.modjam.content.commands.GetAugmentCommand;
+import com.portingdeadmods.modjam.content.commands.SetAugmentCooldownCommand;
+import net.minecraft.commands.CommandSourceStack;
 import net.neoforged.bus.api.SubscribeEvent;
 import net.neoforged.fml.common.EventBusSubscriber;
 import net.neoforged.neoforge.event.RegisterCommandsEvent;
@@ -12,8 +16,11 @@
 public final class MJCommands {
     @SubscribeEvent
     public static void onCommandRegister(RegisterCommandsEvent event) {
-        GetAugmentCommand.register(event.getDispatcher());
-        SetAugmentCommand.register(event.getDispatcher());
-        ConfigCommand.register(event.getDispatcher());
+        CommandDispatcher<CommandSourceStack> dispatcher = event.getDispatcher();
+        GetAugmentCommand.register(dispatcher);
+        SetAugmentCommand.register(dispatcher);
+        GetAugmentCooldownCommand.register(dispatcher);
+        SetAugmentCooldownCommand.register(dispatcher);
+        ConfigCommand.register(dispatcher);
     }
 }

src/main/java/com/portingdeadmods/modjam/utils/AugmentHelper.java

@@ -6,6 +6,7 @@
 import com.portingdeadmods.modjam.api.augments.AugmentType;
 import com.portingdeadmods.modjam.data.MJDataAttachments;
 import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
+import net.minecraft.nbt.CompoundTag;
 import net.minecraft.world.entity.player.Player;
 
 import java.util.HashMap;
@@ -20,9 +21,19 @@ public static Map<AugmentSlot, Augment> getAugments(Player player) {
         return player.getData(MJDataAttachments.AUGMENTS);
     }
 
+    public static Map<AugmentSlot, CompoundTag> getAugmentsData(Player player) {
+        return player.getData(MJDataAttachments.AUGMENTS_EXTRA_DATA);
+    }
+
     public static void setAugment(Player player, AugmentSlot augmentSlot, Augment augment) {
         Map<AugmentSlot, Augment> augments = new HashMap<>(getAugments(player));
         augments.put(augmentSlot, augment);
         player.setData(MJDataAttachments.AUGMENTS, augments);
     }
+
+    public static void setAugmentExtraData(Player player, AugmentSlot augmentSlot, CompoundTag tag) {
+        Map<AugmentSlot, CompoundTag> augments = new HashMap<>(getAugmentsData(player));
+        augments.put(augmentSlot, tag);
+        player.setData(MJDataAttachments.AUGMENTS_EXTRA_DATA, augments);
+    }
 }