augment api partial rewrite

src/main/java/com/portingdeadmods/modjam/MJRegistries.java

@@ -1,15 +1,20 @@
 package com.portingdeadmods.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.api.multiblocks.Multiblock;
-import com.portingdeadmods.modjam.content.augments.StaticAugment;
 import net.minecraft.core.Registry;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.resources.ResourceLocation;
 import net.neoforged.neoforge.registries.RegistryBuilder;
 
 public final class MJRegistries {
-    private static final ResourceKey<Registry<StaticAugment>> AUGMENT_KEY = ResourceKey.createRegistryKey(ResourceLocation.fromNamespaceAndPath(ModJam.MODID, "augment"));
+    private static final ResourceKey<Registry<AugmentType<?>>> AUGMENT_TYPE_KEY = ResourceKey.createRegistryKey(ResourceLocation.fromNamespaceAndPath(ModJam.MODID, "augment_type"));
+    public static final ResourceKey<Registry<AugmentSlot>> AUGMENT_SLOT_KEY = ResourceKey.createRegistryKey(ResourceLocation.fromNamespaceAndPath(ModJam.MODID, "augment_slot"));
     private static final ResourceKey<Registry<Multiblock>> MULTIBLOCK_KEY = ResourceKey.createRegistryKey(ResourceLocation.fromNamespaceAndPath(ModJam.MODID, "multiblock"));
-    public static final Registry<StaticAugment> AUGMENT = new RegistryBuilder<>(AUGMENT_KEY).create();
+
+    public static final Registry<AugmentType<?>> AUGMENT_TYPE = new RegistryBuilder<>(AUGMENT_TYPE_KEY).create();
+    public static final Registry<AugmentSlot> AUGMENT_SLOT = new RegistryBuilder<>(AUGMENT_SLOT_KEY).create();
     public static final Registry<Multiblock> MULTIBLOCK = new RegistryBuilder<>(MULTIBLOCK_KEY).create();
 }

src/main/java/com/portingdeadmods/modjam/ModJam.java

@@ -2,6 +2,7 @@
 
 import com.mojang.logging.LogUtils;
 import com.portingdeadmods.modjam.content.items.PrismMonocleItem;
+import com.portingdeadmods.modjam.data.MJDataAttachments;
 import com.portingdeadmods.modjam.data.MJDataComponents;
 import com.portingdeadmods.modjam.registries.*;
 import net.neoforged.bus.api.IEventBus;
@@ -11,18 +12,16 @@
 import net.neoforged.neoforge.registries.NewRegistryEvent;
 import org.slf4j.Logger;
 
-import java.util.Random;
-
 @Mod(ModJam.MODID)
 public final class ModJam {
     public static final String MODID = "modjam";
     public static final Logger LOGGER = LogUtils.getLogger();
-    public static final Random random = new Random();
 
     public ModJam(IEventBus modEventBus, ModContainer modContainer) {
         modEventBus.addListener(NewRegistryEvent.class, event -> {
             event.register(MJRegistries.MULTIBLOCK);
-            event.register(MJRegistries.AUGMENT);
+            event.register(MJRegistries.AUGMENT_SLOT);
+            event.register(MJRegistries.AUGMENT_TYPE);
         });
 
         MJEntites.ENTITIES.register(modEventBus);
@@ -32,11 +31,13 @@ public ModJam(IEventBus modEventBus, ModContainer modContainer) {
         MJRecipes.SERIALIZERS.register(modEventBus);
         MJFluidTypes.FLUID_TYPES.register(modEventBus);
         MJDataAttachments.ATTACHMENTS.register(modEventBus);
+        MJArgumentTypes.ARGUMENT_TYPES.register(modEventBus);
         MJBlockEntityTypes.BLOCK_ENTITIES.register(modEventBus);
         MJCreativeTabs.CREATIVE_MODE_TABS.register(modEventBus);
         MJDataComponents.DATA_COMPONENT_TYPES.register(modEventBus);
         MJMultiblocks.MULTIBLOCKS.register(modEventBus);
         MJAugments.AUGMENTS.register(modEventBus);
+        MJAugmentSlots.AUGMENT_SLOTS.register(modEventBus);
         MJMenuTypes.MENUS.register(modEventBus);
         MJStructures.STRUCTURES.register(modEventBus);
         MJLootModifier.LOOT_MODIFIERS.register(modEventBus);

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

@@ -0,0 +1,83 @@
+package com.portingdeadmods.modjam.api.augments;
+
+import net.minecraft.core.HolderLookup;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.world.entity.player.Player;
+import net.neoforged.neoforge.common.util.INBTSerializable;
+import net.neoforged.neoforge.event.level.BlockEvent;
+import net.neoforged.neoforge.event.tick.PlayerTickEvent;
+import org.jetbrains.annotations.UnknownNullability;
+
+public abstract class Augment implements INBTSerializable<CompoundTag> {
+    protected final AugmentType<?> augmentType;
+    protected Player player;
+    protected final AugmentSlot augmentSlot;
+
+    // Serialized
+    private int cooldown;
+
+    public Augment(AugmentType<?> augmentType, AugmentSlot augmentSlot) {
+        this.augmentType = augmentType;
+        this.augmentSlot = augmentSlot;
+    }
+
+    public void setPlayer(Player player) {
+        this.player = player;
+    }
+
+    public AugmentType<?> getAugmentType() {
+        return augmentType;
+    }
+
+    public Player getPlayer() {
+        return player;
+    }
+
+    public AugmentSlot getAugmentSlot() {
+        return augmentSlot;
+    }
+
+    public int getCooldown() {
+        return cooldown;
+    }
+
+    public void setCooldown(int cooldown) {
+        this.cooldown = cooldown;
+    }
+
+    public void breakBlock(BlockEvent.BreakEvent event) {
+
+    }
+
+    public void commonTick(PlayerTickEvent.Post event) {
+        if (player.level().isClientSide) clientTick(event);
+        else serverTick(event);
+    }
+
+    @Deprecated
+    public void clientTick(PlayerTickEvent.Post event) {
+
+    }
+
+    @Deprecated
+    public void serverTick(PlayerTickEvent.Post event) {
+
+    }
+
+    public void handleKeybindPress() {
+    }
+
+    public boolean isOnCooldown() {
+        return getCooldown() > 0;
+    }
+
+    @Override
+    public @UnknownNullability CompoundTag serializeNBT(HolderLookup.Provider provider) {
+        return new CompoundTag();
+    }
+
+    @Override
+    public void deserializeNBT(HolderLookup.Provider provider, CompoundTag nbt) {
+
+    }
+}

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

@@ -0,0 +1,13 @@
+package com.portingdeadmods.modjam.api.augments;
+
+import net.minecraft.resources.ResourceLocation;
+
+public interface AugmentSlot {
+    int getSlotId();
+
+    ResourceLocation getLocation();
+
+    default String getName() {
+        return getLocation().getPath();
+    }
+}

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

@@ -0,0 +1,30 @@
+package com.portingdeadmods.modjam.api.augments;
+
+import com.portingdeadmods.modjam.MJRegistries;
+import com.portingdeadmods.modjam.registries.MJAugments;
+
+public class AugmentType<T extends Augment> {
+    private final AugmentConstructor<T> augmentConstructor;
+
+    public static<T extends Augment> AugmentType<T> of(AugmentConstructor<T> constructor) {
+        return new AugmentType<>(constructor);
+    }
+
+    private AugmentType(AugmentConstructor<T> augmentConstructor) {
+        this.augmentConstructor = augmentConstructor;
+    }
+
+    public T create(AugmentSlot augmentSlot) {
+        return augmentConstructor.create(augmentSlot);
+    }
+
+    @Override
+    public String toString() {
+        return MJRegistries.AUGMENT_TYPE.getKey(this).toString();
+    }
+
+    @FunctionalInterface
+    public interface AugmentConstructor<T extends Augment> {
+        T create(AugmentSlot augmentSlot);
+    }
+}

src/main/java/com/portingdeadmods/modjam/client/renderer/augments/AugmentLayerRenderer.java

@@ -2,42 +2,45 @@
 
 import com.mojang.blaze3d.vertex.PoseStack;
 import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.client.model.augment.AugmentModel;
-import com.portingdeadmods.modjam.content.augments.AugmentHelper;
-import com.portingdeadmods.modjam.content.augments.StaticAugment;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.client.model.PlayerModel;
 import net.minecraft.client.player.AbstractClientPlayer;
 import net.minecraft.client.renderer.MultiBufferSource;
 import net.minecraft.client.renderer.RenderType;
 import net.minecraft.client.renderer.entity.RenderLayerParent;
 import net.minecraft.client.renderer.entity.layers.RenderLayer;
-import net.minecraft.client.renderer.entity.player.PlayerRenderer;
 import net.minecraft.client.renderer.texture.OverlayTexture;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.entity.player.Player;
 
+import java.util.Iterator;
+import java.util.List;
+
 public class AugmentLayerRenderer extends RenderLayer<AbstractClientPlayer, PlayerModel<AbstractClientPlayer>> {
     public AugmentLayerRenderer(RenderLayerParent<AbstractClientPlayer, PlayerModel<AbstractClientPlayer>> renderLayerParent) {
         super(renderLayerParent);
     }
 
     @Override
     public void render(PoseStack poseStack, MultiBufferSource bufferSource, int packedLight, AbstractClientPlayer player, float limbSwing, float limbSwingAmount, float partialTick, float ageInTicks, float netHeadYaw, float headPitch) {
-        StaticAugment[] augments = getAugment(player);
-        for (StaticAugment augment : augments) {
+        Iterable<Augment> augments = getAugments(player);
+        for (Augment augment : augments) {
             if (augment != null) {
-                renderAugmentModel(poseStack, bufferSource, packedLight, player, augment.getId());
+                renderAugmentModel(poseStack, bufferSource, packedLight, augment);
             }
         }
     }
 
-    private void renderAugmentModel(PoseStack poseStack, MultiBufferSource bufferSource, int packedLight, Player player, int id) {
+    private void renderAugmentModel(PoseStack poseStack, MultiBufferSource bufferSource, int packedLight, Augment augment) {
         AugmentModel model = new AugmentModel();
         ResourceLocation texture = ResourceLocation.fromNamespaceAndPath(ModJam.MODID, "textures/entity/augment/test_texture.png");
         model.renderToBuffer(poseStack, bufferSource.getBuffer(RenderType.entityCutoutNoCull(texture)), packedLight, OverlayTexture.NO_OVERLAY, 0xFFFFFFFF);
     }
 
-    private StaticAugment[] getAugment(Player player) {
-        return AugmentHelper.getAugments(player);
+    private Iterable<Augment> getAugments(Player player) {
+        return AugmentHelper.getAugments(player).values();
     }
 }

src/main/java/com/portingdeadmods/modjam/content/augments/AugmentSlots.java

@@ -0,0 +1,42 @@
+package com.portingdeadmods.modjam.content.augments;
+
+import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import net.minecraft.resources.ResourceLocation;
+import org.jetbrains.annotations.NotNull;
+
+public enum AugmentSlots implements AugmentSlot {
+    HEAD("head", 0),
+    BODY("body", 1),
+    ARMS("arms", 2),
+    LEGS("legs",3),
+    HEART("heart", 4),
+    NONE("none", -1);
+
+    private final ResourceLocation name;
+    private final int slotId;
+
+    AugmentSlots(String name, int id) {
+        this.name = ResourceLocation.fromNamespaceAndPath(ModJam.MODID, name);
+        this.slotId = id;
+    }
+
+    public static AugmentSlots getValue(int id) {
+        AugmentSlots[] Slots = AugmentSlots.values();
+        for (AugmentSlots slot : Slots) {
+            if (slot.slotId == id)
+                return slot;
+        }
+        return AugmentSlots.NONE;
+    }
+
+    @Override
+    public @NotNull ResourceLocation getLocation() {
+        return name;
+    }
+
+    @Override
+    public int getSlotId() {
+        return slotId;
+    }
+}

src/main/java/com/portingdeadmods/modjam/content/augments/DisallowBreakingAugment.java

@@ -1,16 +1,18 @@
 package com.portingdeadmods.modjam.content.augments;
 
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.registries.MJAugments;
+import net.minecraft.world.entity.player.Player;
 import net.neoforged.neoforge.event.level.BlockEvent;
 
-public class DisallowBreakingAugment extends Augment{
-    @Override
-    public int getId() {
-        return 1;
+public class DisallowBreakingAugment extends Augment {
+    public DisallowBreakingAugment(AugmentSlot augmentSlot) {
+        super(MJAugments.DISALLOW_BREAKING.get(), augmentSlot);
     }
 
     @Override
-    public void breakBlock(Slot slot, BlockEvent.BreakEvent event) {
+    public void breakBlock(BlockEvent.BreakEvent event) {
         event.setCanceled(true);
     }
 }

src/main/java/com/portingdeadmods/modjam/content/augments/GiveDiamondAugment.java

@@ -1,29 +1,30 @@
 package com.portingdeadmods.modjam.content.augments;
 
 import com.mojang.blaze3d.platform.InputConstants;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.network.KeyPressedPayload;
+import com.portingdeadmods.modjam.registries.MJAugments;
 import com.portingdeadmods.modjam.utils.InputUtils;
 import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.item.Items;
 import net.neoforged.neoforge.event.tick.PlayerTickEvent;
 import net.neoforged.neoforge.network.PacketDistributor;
 
 public class GiveDiamondAugment extends Augment {
-    @Override
-    public int getId() {
-        return 2;
+    public GiveDiamondAugment(AugmentSlot augmentSlot) {
+        super(MJAugments.GIVE_DIAMOND.get(), augmentSlot);
     }
 
     @Override
-    public void clientTick(Slot slot, PlayerTickEvent.Post event) {
+    public void clientTick(PlayerTickEvent.Post event) {
         if (InputUtils.isKeyDown(InputConstants.KEY_Y)) {
-            PacketDistributor.sendToServer(new KeyPressedPayload(getId(), slot.slotId));
+            PacketDistributor.sendToServer(new KeyPressedPayload(augmentSlot, augmentSlot.getSlotId()));
         }
     }
 
     @Override
-    public void handleKeybindPress(Slot slot, Player player) {
+    public void handleKeybindPress() {
         player.addItem(Items.DIAMOND.getDefaultInstance());
     }
 }

src/main/java/com/portingdeadmods/modjam/content/augments/PreventPlayerLoseAirAugment.java

@@ -1,17 +1,18 @@
 package com.portingdeadmods.modjam.content.augments;
 
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.registries.MJAugments;
 import net.minecraft.world.entity.player.Player;
 import net.neoforged.neoforge.event.tick.PlayerTickEvent;
 
-public class PreventPlayerLoseAirAugment extends Augment{
-    @Override
-    public int getId() {
-        return 5;
+public class PreventPlayerLoseAirAugment extends Augment {
+    public PreventPlayerLoseAirAugment(AugmentSlot augmentSlot) {
+        super(MJAugments.PREVENT_PLAYER_LOSE_AIR_AUGMENT.get(), augmentSlot);
     }
 
     @Override
-    public void serverTick(Slot slot, PlayerTickEvent.Post event) {
+    public void serverTick(PlayerTickEvent.Post event) {
         Player player = event.getEntity();
         if(player.isUnderWater()){
             player.setAirSupply(player.getMaxAirSupply());

src/main/java/com/portingdeadmods/modjam/content/augments/ThrowBouncingTridentAugment.java

@@ -1,32 +1,34 @@
 package com.portingdeadmods.modjam.content.augments;
 
 import com.mojang.blaze3d.platform.InputConstants;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.content.entites.ThrownBouncingTrident;
 import com.portingdeadmods.modjam.network.KeyPressedPayload;
+import com.portingdeadmods.modjam.registries.MJAugments;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import com.portingdeadmods.modjam.utils.InputUtils;
-import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.item.Items;
 import net.neoforged.neoforge.event.tick.PlayerTickEvent;
 import net.neoforged.neoforge.network.PacketDistributor;
 
-public class ThrowBouncingTridentAugment extends Augment{
-    @Override
-    public int getId() {
-        return 7;
+public class ThrowBouncingTridentAugment extends Augment {
+    public ThrowBouncingTridentAugment(AugmentSlot augmentSlot) {
+        super(MJAugments.THROWN_BOUNCING_TRIDENT_AUGMENT.get(), augmentSlot);
     }
+
     @Override
-    public void clientTick(Slot slot, PlayerTickEvent.Post event) {
-        if (InputUtils.isKeyDown(InputConstants.KEY_Y) && !onCooldown(slot, event.getEntity())) {
-            PacketDistributor.sendToServer(new KeyPressedPayload(getId(), slot.slotId));
+    public void clientTick(PlayerTickEvent.Post event) {
+        if (InputUtils.isKeyDown(InputConstants.KEY_Y) && !isOnCooldown()) {
+            PacketDistributor.sendToServer(new KeyPressedPayload(augmentSlot, augmentSlot.getSlotId()));
         }
     }
 
     @Override
-    public void handleKeybindPress(Slot slot, Player player) {
+    public void handleKeybindPress() {
         ThrownBouncingTrident trident = new ThrownBouncingTrident(player.level(),player,Items.TRIDENT.getDefaultInstance());
         trident.shootFromRotation(player, player.getXRot(), player.getYRot(), 0.0f, 1.5f, 0.0f);
         player.level().addFreshEntity(trident);
-        AugmentHelper.setCooldownAndUpdate(player, slot, 20); // Set the cooldown, which decrements by 1 every tick
+        setCooldown(20); // Set the cooldown, which decrements by 1 every tick
     }
 }

src/main/java/com/portingdeadmods/modjam/content/augments/ThrowRandomPotionAugments.java

@@ -1,12 +1,13 @@
 package com.portingdeadmods.modjam.content.augments;
 
 import com.mojang.blaze3d.platform.InputConstants;
-import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.network.KeyPressedPayload;
+import com.portingdeadmods.modjam.registries.MJAugments;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import com.portingdeadmods.modjam.utils.InputUtils;
 import net.minecraft.core.Holder;
-import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.entity.projectile.ThrownPotion;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.Items;
@@ -19,34 +20,34 @@
 import java.util.ArrayList;
 import java.util.List;
 
-public class ThrowRandomPotionAugments extends Augment{
-    @Override
-    public int getId() {
-        return 4;
+public class ThrowRandomPotionAugments extends Augment {
+    public ThrowRandomPotionAugments(AugmentSlot augmentSlot) {
+        super(MJAugments.THROW_POTION_AUGMENT.get(), augmentSlot);
     }
+
     @Override
-    public void clientTick(Slot slot, PlayerTickEvent.Post event) {
-        if (InputUtils.isKeyDown(InputConstants.KEY_Y) && !onCooldown(slot, event.getEntity())) {
-            PacketDistributor.sendToServer(new KeyPressedPayload(getId(), slot.slotId));
+    public void clientTick(PlayerTickEvent.Post event) {
+        if (InputUtils.isKeyDown(InputConstants.KEY_Y) && !isOnCooldown()) {
+            PacketDistributor.sendToServer(new KeyPressedPayload(augmentSlot, augmentSlot.getSlotId()));
         }
     }
 
     @Override
-    public void handleKeybindPress(Slot slot, Player player) {
+    public void handleKeybindPress() {
         List<Holder<Potion>> potions = new ArrayList<>();
         potions.add(Potions.HEALING);
         potions.add(Potions.HARMING);
         potions.add(Potions.REGENERATION);
         potions.add(Potions.SWIFTNESS);
         potions.add(Potions.SLOWNESS);
 
-        Holder<Potion> randomPotion = potions.get(ModJam.random.nextInt(potions.size()));
+        Holder<Potion> randomPotion = potions.get(player.getRandom().nextInt(potions.size()));
         ItemStack stack = PotionContents.createItemStack(Items.SPLASH_POTION,randomPotion);
 
         ThrownPotion potion = new ThrownPotion(player.level(),player);
         potion.setItem(stack);
         potion.shootFromRotation(player, player.getXRot(), player.getYRot(), 0.0F, 1.5F, 1.0F);
         player.level().addFreshEntity(potion);
-        AugmentHelper.setCooldownAndUpdate(player, slot, 20); // Set the cooldown, which decrements by 1 every tick
+        setCooldown(20); // Set the cooldown, which decrements by 1 every tick
     }
 }

src/main/java/com/portingdeadmods/modjam/content/augments/ThrowSnowballAugment.java

@@ -1,8 +1,11 @@
 package com.portingdeadmods.modjam.content.augments;
 
 import com.mojang.blaze3d.platform.InputConstants;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
 import com.portingdeadmods.modjam.network.KeyPressedPayload;
+import com.portingdeadmods.modjam.registries.MJAugments;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import com.portingdeadmods.modjam.utils.InputUtils;
 import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.entity.projectile.Snowball;
@@ -11,24 +14,23 @@
 import net.neoforged.neoforge.network.PacketDistributor;
 
 public class ThrowSnowballAugment extends Augment {
-    @Override
-    public int getId() {
-        return 3;
+    public ThrowSnowballAugment(AugmentSlot augmentSlot) {
+        super(MJAugments.THROW_SNOWBALL.get(), augmentSlot);
     }
 
     @Override
-    public void clientTick(Slot slot, PlayerTickEvent.Post event) {
-        if (InputUtils.isKeyDown(InputConstants.KEY_Y) && !onCooldown(slot, event.getEntity())) {
-            PacketDistributor.sendToServer(new KeyPressedPayload(getId(), slot.slotId));
+    public void clientTick(PlayerTickEvent.Post event) {
+        if (InputUtils.isKeyDown(InputConstants.KEY_Y) && !isOnCooldown()) {
+            PacketDistributor.sendToServer(new KeyPressedPayload(augmentSlot, augmentSlot.getSlotId()));
         }
     }
 
     @Override
-    public void handleKeybindPress(Slot slot, Player player) {
+    public void handleKeybindPress() {
         Snowball snowball = new Snowball(player.level(), player);
         snowball.setItem(Items.SNOWBALL.getDefaultInstance());
         snowball.shootFromRotation(player, player.getXRot(), player.getYRot(), 0.0F, 1.5F, 1.0F);
         player.level().addFreshEntity(snowball);
-        AugmentHelper.setCooldownAndUpdate(player, slot, 20); // Set the cooldown, which decrements by 1 every tick
+        setCooldown(20); // Set the cooldown, which decrements by 1 every tick
     }
 }

src/main/java/com/portingdeadmods/modjam/content/augments/UnderwaterMovementSpeedAugment.java

@@ -1,18 +1,21 @@
 package com.portingdeadmods.modjam.content.augments;
 
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.registries.MJAugments;
 import net.minecraft.world.effect.MobEffectInstance;
 import net.minecraft.world.effect.MobEffects;
+import net.minecraft.world.entity.player.Player;
+import net.minecraft.world.inventory.Slot;
 import net.neoforged.neoforge.event.tick.PlayerTickEvent;
 
-public class UnderwaterMovementSpeedAugment extends Augment{
-    @Override
-    public int getId() {
-        return 6;
+public class UnderwaterMovementSpeedAugment extends Augment {
+    public UnderwaterMovementSpeedAugment(AugmentSlot augmentSlot) {
+        super(MJAugments.UNDERWATER_MOVEMENT_SPEED_AUGMENT.get(), augmentSlot);
     }
 
     @Override
-    public void serverTick(Slot slot, PlayerTickEvent.Post event) {
+    public void serverTick(PlayerTickEvent.Post event) {
         if (event.getEntity().isUnderWater()){
             event.getEntity().addEffect(new MobEffectInstance(MobEffects.DOLPHINS_GRACE,20,1));
         }

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

@@ -0,0 +1,52 @@
+package com.portingdeadmods.modjam.content.commands;
+
+import com.mojang.brigadier.StringReader;
+import com.mojang.brigadier.arguments.ArgumentType;
+import com.mojang.brigadier.context.CommandContext;
+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.AugmentSlot;
+import com.portingdeadmods.modjam.api.augments.AugmentType;
+import net.minecraft.commands.SharedSuggestionProvider;
+import net.minecraft.network.chat.Component;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.resources.ResourceLocation;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+public class AugmentSlotArgumentType implements ArgumentType<AugmentSlot> {
+    private static final AugmentSlotArgumentType INSTANCE = new AugmentSlotArgumentType();
+
+    private static final DynamicCommandExceptionType UNKNOWN_TYPE = new DynamicCommandExceptionType(
+            type -> Component.literal("Unknown augment slot: " + type));
+    public static Set<String> suggestions = new HashSet<>();
+
+    private AugmentSlotArgumentType() {
+    }
+
+    public static AugmentSlotArgumentType getInstance() {
+        return INSTANCE;
+    }
+
+    @Override
+    public AugmentSlot parse(StringReader reader) throws CommandSyntaxException {
+        ResourceLocation read = ResourceLocation.read(reader);
+        ModJam.LOGGER.debug("Res: {}", read);
+        AugmentSlot augmentSlot = MJRegistries.AUGMENT_SLOT.get(ResourceKey.create(MJRegistries.AUGMENT_SLOT_KEY, read));
+        if (augmentSlot != null) {
+            return augmentSlot;
+        }
+        throw UNKNOWN_TYPE.create(read.toString());
+    }
+
+    @Override
+    public <S> CompletableFuture<Suggestions> listSuggestions(CommandContext<S> context, SuggestionsBuilder builder) {
+        return SharedSuggestionProvider.suggest(suggestions, builder);
+    }
+}

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

@@ -0,0 +1,50 @@
+package com.portingdeadmods.modjam.content.commands;
+
+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;
+import net.minecraft.resources.ResourceLocation;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+public class AugmentTypeArgumentType implements ArgumentType<AugmentType<?>> {
+    private static final AugmentTypeArgumentType INSTANCE = new AugmentTypeArgumentType();
+
+    private static final DynamicCommandExceptionType UNKNOWN_TYPE = new DynamicCommandExceptionType(
+            type -> Component.literal("Unknown augment type: " + type));
+    public static Set<String> suggestions = new HashSet<>();
+
+    private AugmentTypeArgumentType() {
+    }
+
+    public static AugmentTypeArgumentType getInstance() {
+        return INSTANCE;
+    }
+
+    @Override
+    public <S> CompletableFuture<Suggestions> listSuggestions(CommandContext<S> context, SuggestionsBuilder builder) {
+        return SharedSuggestionProvider.suggest(suggestions, builder);
+    }
+
+    @Override
+    public AugmentType<?> parse(StringReader reader) throws CommandSyntaxException {
+        ResourceLocation read = ResourceLocation.read(reader);
+        AugmentType<?> augmentType = MJRegistries.AUGMENT_TYPE.get(read);
+        if (augmentType != null) {
+            return augmentType;
+        }
+        throw UNKNOWN_TYPE.create(read.toString());
+    }
+}

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

@@ -0,0 +1,39 @@
+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.augments.AugmentSlots;
+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>
+public class GetAugmentCommand {
+    public static void register(CommandDispatcher<CommandSourceStack> dispatcher) {
+        dispatcher.register(Commands.literal(ModJam.MODID)
+                .then(Commands.literal("augments")
+                        .then(Commands.literal("get")
+                                .then(Commands.argument("slot", AugmentSlotArgumentType.getInstance())
+                                        .executes(GetAugmentCommand::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);
+        String augment = "None";
+        if (augmentBySlot != null) {
+            AugmentType<?> augmentType = augmentBySlot.getAugmentType();
+            augment = augmentType.toString();
+        }
+        player.sendSystemMessage(Component.literal("Augment in slot '" + augmentSlot.getName() + "': " + augment));
+        return 1;
+    }
+}

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

@@ -0,0 +1,42 @@
+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.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>
+
+// TODO: Only set augments for slots that support them
+public class SetAugmentCommand {
+    public static void register(CommandDispatcher<CommandSourceStack> dispatcher) {
+        dispatcher.register(Commands.literal(ModJam.MODID)
+                .then(Commands.literal("augments")
+                        .then(Commands.literal("set")
+                                .then(Commands.argument("slot", AugmentSlotArgumentType.getInstance())
+                                        .then(Commands.argument("augment", AugmentTypeArgumentType.getInstance())
+                                                .executes(SetAugmentCommand::execute))))));
+    }
+
+    private static int execute(CommandContext<CommandSourceStack> ctx) {
+        Player player = ctx.getSource().getPlayer();
+        AugmentSlot slot = ctx.getArgument("slot", AugmentSlot.class);
+        AugmentType<?> type = ctx.getArgument("augment", AugmentType.class);
+        Augment augment = type.create(slot);
+        augment.setPlayer(player);
+        AugmentHelper.setAugment(player, slot, augment);
+        player.sendSystemMessage(Component.literal("Set augment in slot '" + slot.getName() + "' to: " + type));
+        return 1;
+    }
+
+}

src/main/java/com/portingdeadmods/modjam/content/multiblocks/AugmentationStationMultiblock.java

@@ -0,0 +1,51 @@
+package com.portingdeadmods.modjam.content.multiblocks;
+
+import com.portingdeadmods.modjam.api.blockentities.multiblock.MultiblockEntity;
+import com.portingdeadmods.modjam.api.multiblocks.Multiblock;
+import com.portingdeadmods.modjam.api.multiblocks.MultiblockData;
+import com.portingdeadmods.modjam.api.multiblocks.MultiblockLayer;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import net.minecraft.core.BlockPos;
+import net.minecraft.world.entity.player.Player;
+import net.minecraft.world.level.Level;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.entity.BlockEntityType;
+import net.minecraft.world.level.block.state.BlockState;
+import org.jetbrains.annotations.Nullable;
+
+public class AugmentationStationMultiblock implements Multiblock {
+    @Override
+    public Block getUnformedController() {
+        return null;
+    }
+
+    @Override
+    public Block getFormedController() {
+        return null;
+    }
+
+    @Override
+    public MultiblockLayer[] getLayout() {
+        return new MultiblockLayer[0];
+    }
+
+    @Override
+    public Int2ObjectMap<Block> getDefinition() {
+        return null;
+    }
+
+    @Override
+    public BlockEntityType<? extends MultiblockEntity> getMultiBlockEntityType() {
+        return null;
+    }
+
+    @Override
+    public @Nullable BlockState formBlock(Level level, BlockPos blockPos, BlockPos controllerPos, int layerIndex, int layoutIndex, MultiblockData multiblockData, @Nullable Player player) {
+        return null;
+    }
+
+    @Override
+    public boolean isFormed(Level level, BlockPos blockPos) {
+        return false;
+    }
+}

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

@@ -0,0 +1,28 @@
+package com.portingdeadmods.modjam.data;
+
+import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.utils.AugmentCodecs;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
+import net.minecraft.nbt.CompoundTag;
+import net.neoforged.neoforge.attachment.AttachmentType;
+import net.neoforged.neoforge.registries.DeferredRegister;
+import net.neoforged.neoforge.registries.NeoForgeRegistries;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.function.Supplier;
+
+public final class MJDataAttachments {
+    public static final DeferredRegister<AttachmentType<?>> ATTACHMENTS = DeferredRegister.create(NeoForgeRegistries.ATTACHMENT_TYPES, ModJam.MODID);
+
+    public static final Supplier<AttachmentType<Map<AugmentSlot, Augment>>> AUGMENTS = ATTACHMENTS.register(
+            "augments", () -> AttachmentType.<Map<AugmentSlot, Augment>>builder(Collections::emptyMap)
+                    .serialize(AugmentCodecs.AUGMENTS_CODEC).copyOnDeath().build()
+    );
+    public static final Supplier<AttachmentType<Map<AugmentSlot, CompoundTag>>> AUGMENTS_EXTRA_DATA = ATTACHMENTS.register(
+            "augments_extra_data", () -> AttachmentType.<Map<AugmentSlot, CompoundTag>>builder(Collections::emptyMap)
+                    .serialize(AugmentCodecs.AUGMENTS_EXTRA_DATA_CODEC).copyOnDeath().build()
+    );
+}

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

@@ -1,46 +1,50 @@
 package com.portingdeadmods.modjam.events;
 
 import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
-import com.portingdeadmods.modjam.content.augments.AugmentHelper;
-import com.portingdeadmods.modjam.content.augments.StaticAugment;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.world.entity.player.Player;
 import net.neoforged.bus.api.SubscribeEvent;
 import net.neoforged.fml.common.EventBusSubscriber;
 import net.neoforged.neoforge.event.level.BlockEvent;
 import net.neoforged.neoforge.event.tick.PlayerTickEvent;
 
+import java.util.List;
+
 @SuppressWarnings("unused")
 @EventBusSubscriber(modid = ModJam.MODID)
 public class AugmentEvents {
 
     @SubscribeEvent
     public static void breakEvent(BlockEvent.BreakEvent event){
-        StaticAugment[] augments = AugmentHelper.getAugments(event.getPlayer());
-        for (int i = 0; i < augments.length; i++) {
-            if (augments[i] != null){
-                augments[i].breakBlock(Slot.GetValue(i),event);
+        Iterable<Augment> augments = AugmentHelper.getAugments(event.getPlayer()).values();
+        for (Augment augment : augments) {
+            if (augment != null) {
+                //augments.get(i).breakBlock(AugmentSlot.GetValue(i),event);
+
             }
         }
     }
 
 
     @SubscribeEvent
     public static void playerTick(PlayerTickEvent.Post event){
-        StaticAugment[] augments = AugmentHelper.getAugments(event.getEntity());
-        for (int i = 0; i < 5; i++) {
-            StaticAugment augment = augments[i];
+        Iterable<Augment> augments = AugmentHelper.getAugments(event.getEntity()).values();
+        for (Augment augment : augments) {
             if (augment != null) {
-                Slot slot = Slot.GetValue(i);
+                AugmentSlot slot = augment.getAugmentSlot();
                 Player player = event.getEntity();
 
                 if (player.level().isClientSide) {
-                    augment.clientTick(slot,event);
-                    if (AugmentHelper.getCooldown(player, slot) >= 0){
-                        AugmentHelper.setCooldownAndUpdate(player, slot, AugmentHelper.getCooldown(player, slot) - 1);
+                    augment.clientTick(event);
+                    // TODO: Move this to tick method of Augment class
+                    if (augment.getCooldown() >= 0){
+                        augment.setCooldown(augment.getCooldown() - 1);
                     }
                 } else {
-                    augment.serverTick(slot,event);
+                    augment.serverTick(event);
                 }
             }
         }

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

@@ -1,10 +1,17 @@
 package com.portingdeadmods.modjam.events;
 
+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 com.portingdeadmods.modjam.content.commands.AugmentSlotArgumentType;
+import com.portingdeadmods.modjam.content.commands.AugmentTypeArgumentType;
 import com.portingdeadmods.modjam.content.recipes.ItemEtchingRecipe;
 import com.portingdeadmods.modjam.registries.MJFluidTypes;
 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.world.entity.item.ItemEntity;
 import net.minecraft.world.item.ItemStack;
@@ -15,65 +22,86 @@
 import net.neoforged.bus.api.SubscribeEvent;
 import net.neoforged.fml.common.EventBusSubscriber;
 import net.neoforged.neoforge.event.tick.EntityTickEvent;
+import net.neoforged.neoforge.registries.RegisterEvent;
 
+import java.util.Objects;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
-@EventBusSubscriber(modid = "modjam",bus = EventBusSubscriber.Bus.GAME)
 public final class MJEvents {
 
-    private static final Object2IntMap<ItemEntity> activeTransformations = new Object2IntOpenHashMap<>();
+    @EventBusSubscriber(modid = ModJam.MODID, bus = EventBusSubscriber.Bus.GAME)
+    public static class Game {
+        private static final Object2IntMap<ItemEntity> activeTransformations = new Object2IntOpenHashMap<>();
 
-    @SubscribeEvent
-    public static void onItemEntityTick(EntityTickEvent.Post event) {
-        if (event.getEntity() instanceof ItemEntity itemEntity) {
-            Level level = itemEntity.level();
+        @SubscribeEvent
+        public static void onItemEntityTick(EntityTickEvent.Post event) {
+            if (event.getEntity() instanceof ItemEntity itemEntity) {
+                Level level = itemEntity.level();
 
-            if (itemEntity.isInFluidType(MJFluidTypes.ETCHING_ACID_FLUID_TYPE.get())) {
-                processItemEtching(itemEntity, level);
+                if (itemEntity.isInFluidType(MJFluidTypes.ETCHING_ACID_FLUID_TYPE.get())) {
+                    processItemEtching(itemEntity, level);
+                }
             }
         }
-    }
 
-    private static void processItemEtching(ItemEntity itemEntity, Level level) {
-        ItemStack stack = itemEntity.getItem();
-
-        if (!activeTransformations.containsKey(itemEntity)) {
-            Optional<ItemEtchingRecipe> optionalRecipe = getEtchingRecipe(stack, level);
-            if (optionalRecipe.isPresent()) {
-                activeTransformations.put(itemEntity, 0);
-            }
-        } else {
-            int etchingTime = activeTransformations.getInt(itemEntity);
+        private static void processItemEtching(ItemEntity itemEntity, Level level) {
+            ItemStack stack = itemEntity.getItem();
 
-            if (etchingTime >= 100) {
+            if (!activeTransformations.containsKey(itemEntity)) {
                 Optional<ItemEtchingRecipe> optionalRecipe = getEtchingRecipe(stack, level);
                 if (optionalRecipe.isPresent()) {
-                    transformItem(itemEntity, optionalRecipe.get(), level);
+                    activeTransformations.put(itemEntity, 0);
                 }
-                activeTransformations.remove(itemEntity);
             } else {
-                activeTransformations.put(itemEntity, etchingTime + 1);
-                // Optionally spawn particles while etching
-                ParticlesUtils.spawnParticles(itemEntity, level, ParticleTypes.FLAME);
+                int etchingTime = activeTransformations.getInt(itemEntity);
+
+                if (etchingTime >= 100) {
+                    Optional<ItemEtchingRecipe> optionalRecipe = getEtchingRecipe(stack, level);
+                    if (optionalRecipe.isPresent()) {
+                        transformItem(itemEntity, optionalRecipe.get(), level);
+                    }
+                    activeTransformations.remove(itemEntity);
+                } else {
+                    activeTransformations.put(itemEntity, etchingTime + 1);
+                    // Optionally spawn particles while etching
+                    ParticlesUtils.spawnParticles(itemEntity, level, ParticleTypes.FLAME);
+                }
             }
         }
-    }
 
-    private static Optional<ItemEtchingRecipe> getEtchingRecipe(ItemStack stack, Level level) {
-        return level.getRecipeManager()
-                .getRecipeFor(ItemEtchingRecipe.Type.INSTANCE, new SingleRecipeInput(stack), level)
-                .map(RecipeHolder::value);
-    }
+        private static Optional<ItemEtchingRecipe> getEtchingRecipe(ItemStack stack, Level level) {
+            return level.getRecipeManager()
+                    .getRecipeFor(ItemEtchingRecipe.Type.INSTANCE, new SingleRecipeInput(stack), level)
+                    .map(RecipeHolder::value);
+        }
+
+        private static void transformItem(ItemEntity itemEntity, ItemEtchingRecipe recipe, Level level) {
+            Vec3 position = itemEntity.position();
 
-    private static void transformItem(ItemEntity itemEntity, ItemEtchingRecipe recipe, Level level) {
-        Vec3 position = itemEntity.position();
+            ItemStack resultStack = recipe.getResultItem(level.registryAccess()).copy();
+            resultStack.setCount(itemEntity.getItem().getCount());
 
-        ItemStack resultStack = recipe.getResultItem(level.registryAccess()).copy();
-        resultStack.setCount(itemEntity.getItem().getCount());
+            itemEntity.discard();
+
+            ItemEntity newItemEntity = new ItemEntity(level, position.x, position.y, position.z, resultStack);
+            level.addFreshEntity(newItemEntity);
+        }
+    }
 
-        itemEntity.discard();
+    @EventBusSubscriber(modid = ModJam.MODID, bus = EventBusSubscriber.Bus.MOD)
+    public static class Mod {
+        @SubscribeEvent
+        public static void onRegisterAugments(RegisterEvent event) {
+            Registry<AugmentSlot> slotRegistry = event.getRegistry(MJRegistries.AUGMENT_SLOT.key());
+            if (slotRegistry != null) {
+                AugmentSlotArgumentType.suggestions = slotRegistry.keySet().stream().map(Objects::toString).collect(Collectors.toSet());
+            }
 
-        ItemEntity newItemEntity = new ItemEntity(level, position.x, position.y, position.z, resultStack);
-        level.addFreshEntity(newItemEntity);
+            Registry<AugmentType<?>> augmentRegistry = event.getRegistry(MJRegistries.AUGMENT_TYPE.key());
+            if (augmentRegistry != null) {
+                AugmentTypeArgumentType.suggestions = augmentRegistry.keySet().stream().map(Objects::toString).collect(Collectors.toSet());
+            }
+        }
     }
 }

src/main/java/com/portingdeadmods/modjam/network/KeyPressedPayload.java

@@ -1,8 +1,10 @@
 package com.portingdeadmods.modjam.network;
 
 import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
-import com.portingdeadmods.modjam.content.augments.AugmentHelper;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import com.portingdeadmods.modjam.utils.AugmentCodecs;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.network.RegistryFriendlyByteBuf;
 import net.minecraft.network.chat.Component;
 import net.minecraft.network.codec.ByteBufCodecs;
@@ -11,29 +13,29 @@
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.entity.player.Player;
 import net.neoforged.neoforge.network.handling.IPayloadContext;
+import org.jetbrains.annotations.NotNull;
 
-public record KeyPressedPayload(int augmentId, int slot) implements CustomPacketPayload {
+public record KeyPressedPayload(AugmentSlot augmentSlot, int slot) implements CustomPacketPayload {
     public static final Type<KeyPressedPayload> TYPE = new Type<>(ResourceLocation.fromNamespaceAndPath(ModJam.MODID, "key_pressesd_paylad"));
     public static final StreamCodec<RegistryFriendlyByteBuf, KeyPressedPayload> STREAM_CODEC = StreamCodec.composite(
-            ByteBufCodecs.INT,
-            KeyPressedPayload::augmentId,
+            AugmentCodecs.AUGMENT_SLOT_STREAM_CODEC,
+            KeyPressedPayload::augmentSlot,
             ByteBufCodecs.INT,
             KeyPressedPayload::slot,
             KeyPressedPayload::new
     );
 
     @Override
-    public Type<? extends CustomPacketPayload> type() {
+    public @NotNull Type<? extends CustomPacketPayload> type() {
         return TYPE;
     }
 
-    public static void keyPressedAction(KeyPressedPayload payload, IPayloadContext context){
-        context.enqueueWork(()->{
+    public static void keyPressedAction(KeyPressedPayload payload, IPayloadContext context) {
+        context.enqueueWork(() -> {
             Player player = context.player();
-            int augmentId = payload.augmentId();
-            AugmentHelper.getAugment(augmentId).handleKeybindPress(Slot.GetValue(payload.slot()), player);
-        }).exceptionally(e->{
-            context.disconnect(Component.literal("action failed:  "+ e.getMessage()));
+            AugmentHelper.getAugmentBySlot(player, payload.augmentSlot).handleKeybindPress();
+        }).exceptionally(e -> {
+            context.disconnect(Component.literal("action failed:  " + e.getMessage()));
             return null;
         });
 

src/main/java/com/portingdeadmods/modjam/network/SetAugmentDataPayload.java

@@ -1,8 +1,10 @@
 package com.portingdeadmods.modjam.network;
 
 import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
-import com.portingdeadmods.modjam.content.augments.AugmentHelper;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.network.RegistryFriendlyByteBuf;
 import net.minecraft.network.codec.ByteBufCodecs;
 import net.minecraft.network.codec.StreamCodec;
@@ -27,8 +29,6 @@ public record SetAugmentDataPayload(int augmentId, int slot) implements CustomPa
         return TYPE;
     }
     public static void setAugmentDataAction(SetAugmentDataPayload payload, IPayloadContext context){
-        context.enqueueWork(()->{
-            AugmentHelper.setId(context.player(), Slot.GetValue(payload.slot()), payload.augmentId());
-        });
+//        context.enqueueWork(()-> AugmentHelper.setId(context.player(), AugmentSlots.getValue(payload.slot()), payload.augmentId()));
     }
 }

src/main/java/com/portingdeadmods/modjam/network/SetCooldownPayload.java

@@ -1,8 +1,9 @@
 package com.portingdeadmods.modjam.network;
 
 import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.capabilities.augmentation.Slot;
-import com.portingdeadmods.modjam.content.augments.AugmentHelper;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import com.portingdeadmods.modjam.utils.AugmentHelper;
 import net.minecraft.network.RegistryFriendlyByteBuf;
 import net.minecraft.network.codec.ByteBufCodecs;
 import net.minecraft.network.codec.StreamCodec;
@@ -27,8 +28,6 @@ public record SetCooldownPayload (int cooldown, int slot)implements CustomPacket
     }
 
     public static void setCooldownAction(SetCooldownPayload payload, IPayloadContext context){
-        context.enqueueWork(()->{
-            AugmentHelper.setCooldown(context.player(), Slot.GetValue(payload.slot()), payload.cooldown());
-        });
+//        context.enqueueWork(()-> AugmentHelper.setCooldown(context.player(), AugmentSlots.getValue(payload.slot()), payload.cooldown()));
     }
 }

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

@@ -0,0 +1,28 @@
+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 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;
+
+public final class MJArgumentTypes {
+    public static final DeferredRegister<ArgumentTypeInfo<?, ?>> ARGUMENT_TYPES =
+            DeferredRegister.create(Registries.COMMAND_ARGUMENT_TYPE, ModJam.MODID);
+
+    public static final Supplier<ArgumentTypeInfo<?, ?>> AUGMENT_SLOT_ARGUMENT =
+            ARGUMENT_TYPES.register("augment_slot",
+                    () -> ArgumentTypeInfos.registerByClass(AugmentSlotArgumentType.class,
+                            SingletonArgumentInfo.contextFree(AugmentSlotArgumentType::getInstance)));
+    public static final Supplier<ArgumentTypeInfo<?, ?>> AUGMENT_TYPE_ARGUMENT =
+            ARGUMENT_TYPES.register("augment_type",
+                    () -> ArgumentTypeInfos.registerByClass(AugmentTypeArgumentType.class,
+                            SingletonArgumentInfo.contextFree(AugmentTypeArgumentType::getInstance)));
+}

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

@@ -0,0 +1,19 @@
+package com.portingdeadmods.modjam.registries;
+
+import com.portingdeadmods.modjam.MJRegistries;
+import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.content.augments.AugmentSlots;
+import net.neoforged.neoforge.registries.DeferredRegister;
+
+import java.util.function.Supplier;
+
+public final class MJAugmentSlots {
+    public static final DeferredRegister<AugmentSlot> AUGMENT_SLOTS = DeferredRegister.create(MJRegistries.AUGMENT_SLOT, ModJam.MODID);
+
+    public static final Supplier<AugmentSlot> HEAD = AUGMENT_SLOTS.register("head", () -> AugmentSlots.HEAD);
+    public static final Supplier<AugmentSlot> BODY = AUGMENT_SLOTS.register("body", () -> AugmentSlots.BODY);
+    public static final Supplier<AugmentSlot> ARMS = AUGMENT_SLOTS.register("arms", () -> AugmentSlots.ARMS);
+    public static final Supplier<AugmentSlot> LEGS = AUGMENT_SLOTS.register("legs", () -> AugmentSlots.LEGS);
+    public static final Supplier<AugmentSlot> HEART = AUGMENT_SLOTS.register("heart", () -> AugmentSlots.HEART);
+}

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

@@ -2,21 +2,28 @@
 
 import com.portingdeadmods.modjam.MJRegistries;
 import com.portingdeadmods.modjam.ModJam;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentType;
 import com.portingdeadmods.modjam.content.augments.*;
 import net.neoforged.neoforge.registries.DeferredRegister;
 
 import java.util.function.Supplier;
 
 public final class MJAugments {
-    public static final DeferredRegister<StaticAugment> AUGMENTS = DeferredRegister.create(MJRegistries.AUGMENT, ModJam.MODID);
+    public static final DeferredRegister<AugmentType<?>> AUGMENTS = DeferredRegister.create(MJRegistries.AUGMENT_TYPE, ModJam.MODID);
 
-
-    public static final Supplier<EmptyAugment> EMPTY_AUGMENT = AUGMENTS.register("empty_augment", EmptyAugment::new);
-    public static final Supplier<DisallowBreakingAugment> DISALLOW_BREAKING = AUGMENTS.register("disallow_breaking", DisallowBreakingAugment::new);
-    public static final Supplier<GiveDiamondAugment> GIVE_DIAMOND = AUGMENTS.register("give_diamond", GiveDiamondAugment::new);
-    public static final Supplier<ThrowSnowballAugment> THROW_SNOWBALL = AUGMENTS.register("throw_snowball", ThrowSnowballAugment::new);
-    public static final Supplier<ThrowRandomPotionAugments> THROW_POTION_AUGMENT = AUGMENTS.register("throw_random_potion", ThrowRandomPotionAugments::new);
-    public static final Supplier<PreventPlayerLoseAirAugment> PREVENT_PLAYER_LOSE_AIR_AUGMENT = AUGMENTS.register("prevent_player_lose_air_supply", PreventPlayerLoseAirAugment::new);
-    public static final Supplier<UnderwaterMovementSpeedAugment> UNDERWATER_MOVEMENT_SPEED_AUGMENT = AUGMENTS.register("underwater_movement_speed", UnderwaterMovementSpeedAugment::new);
-    public static final Supplier<ThrowBouncingTridentAugment> THROWN_BOUNCING_TRIDENT_AUGMENT = AUGMENTS.register("throw_bouncing_trident", ThrowBouncingTridentAugment::new);
+    public static final Supplier<AugmentType<DisallowBreakingAugment>> DISALLOW_BREAKING = AUGMENTS.register("disallow_breaking",
+            () -> AugmentType.of(DisallowBreakingAugment::new));
+    public static final Supplier<AugmentType<GiveDiamondAugment>> GIVE_DIAMOND = AUGMENTS.register("give_diamond",
+            () -> AugmentType.of(GiveDiamondAugment::new));
+    public static final Supplier<AugmentType<ThrowSnowballAugment>> THROW_SNOWBALL = AUGMENTS.register("throw_snowball",
+            () -> AugmentType.of(ThrowSnowballAugment::new));
+    public static final Supplier<AugmentType<ThrowRandomPotionAugments>> THROW_POTION_AUGMENT = AUGMENTS.register("throw_random_potion",
+            () -> AugmentType.of(ThrowRandomPotionAugments::new));
+    public static final Supplier<AugmentType<PreventPlayerLoseAirAugment>> PREVENT_PLAYER_LOSE_AIR_AUGMENT = AUGMENTS.register("prevent_player_lose_air_supply",
+            () -> AugmentType.of(PreventPlayerLoseAirAugment::new));
+    public static final Supplier<AugmentType<UnderwaterMovementSpeedAugment>> UNDERWATER_MOVEMENT_SPEED_AUGMENT = AUGMENTS.register("underwater_movement_speed",
+            () -> AugmentType.of(UnderwaterMovementSpeedAugment::new));
+    public static final Supplier<AugmentType<ThrowBouncingTridentAugment>> THROWN_BOUNCING_TRIDENT_AUGMENT = AUGMENTS.register("throw_bouncing_trident",
+            () -> AugmentType.of(ThrowBouncingTridentAugment::new));
 }

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

@@ -1,9 +1,8 @@
 package com.portingdeadmods.modjam.registries;
 
 import com.portingdeadmods.modjam.ModJam;
-import com.portingdeadmods.modjam.content.commands.SetAugmentIdCommand;
-import com.portingdeadmods.modjam.content.commands.ShowAugmentIdCommand;
-import com.portingdeadmods.modjam.content.commands.TestCommand;
+import com.portingdeadmods.modjam.content.commands.SetAugmentCommand;
+import com.portingdeadmods.modjam.content.commands.GetAugmentCommand;
 import net.neoforged.bus.api.SubscribeEvent;
 import net.neoforged.fml.common.EventBusSubscriber;
 import net.neoforged.neoforge.event.RegisterCommandsEvent;
@@ -13,9 +12,8 @@
 public final class MJCommands {
     @SubscribeEvent
     public static void onCommandRegister(RegisterCommandsEvent event) {
-        TestCommand.register(event.getDispatcher());
-        ShowAugmentIdCommand.register(event.getDispatcher());
-        SetAugmentIdCommand.register(event.getDispatcher());
+        GetAugmentCommand.register(event.getDispatcher());
+        SetAugmentCommand.register(event.getDispatcher());
         ConfigCommand.register(event.getDispatcher());
     }
 }

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

@@ -23,7 +23,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
 
-public class ArmorModelsHandler {
+public final class ArmorModelsHandler {
 
     private static final Map<ModelLayerLocation, Layer> LAYERS = new HashMap<>();
     private static final Map<Pair<ModelLayerLocation, EquipmentSlot>, MJArmorModel> CACHED_ARMORS = new HashMap<>();

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

@@ -0,0 +1,31 @@
+package com.portingdeadmods.modjam.utils;
+
+import com.mojang.serialization.Codec;
+import com.mojang.serialization.codecs.RecordCodecBuilder;
+import com.portingdeadmods.modjam.MJRegistries;
+import com.portingdeadmods.modjam.api.augments.Augment;
+import com.portingdeadmods.modjam.api.augments.AugmentSlot;
+import com.portingdeadmods.modjam.api.augments.AugmentType;
+import io.netty.buffer.ByteBuf;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.network.codec.ByteBufCodecs;
+import net.minecraft.network.codec.StreamCodec;
+import net.minecraft.resources.ResourceLocation;
+
+import java.util.Map;
+
+public final class AugmentCodecs {
+    public static final Codec<AugmentSlot> AUGMENT_SLOT_CODEC =
+            ResourceLocation.CODEC.xmap(MJRegistries.AUGMENT_SLOT::get, MJRegistries.AUGMENT_SLOT::getKey);
+    public static final Codec<AugmentType<?>> AUGMENT_TYPE_CODEC =
+            ResourceLocation.CODEC.xmap(MJRegistries.AUGMENT_TYPE::get, MJRegistries.AUGMENT_TYPE::getKey);
+    public static final Codec<Augment> AUGMENT_CODEC = RecordCodecBuilder.create(builder -> builder.group(
+            AUGMENT_TYPE_CODEC.fieldOf("type").forGetter(Augment::getAugmentType),
+            AUGMENT_SLOT_CODEC.fieldOf("slot").forGetter(Augment::getAugmentSlot)
+    ).apply(builder, AugmentType::create));
+    public static final Codec<Map<AugmentSlot, Augment>> AUGMENTS_CODEC = Codec.unboundedMap(AUGMENT_SLOT_CODEC, AUGMENT_CODEC);
+    public static final Codec<Map<AugmentSlot, CompoundTag>> AUGMENTS_EXTRA_DATA_CODEC = Codec.unboundedMap(AUGMENT_SLOT_CODEC, CompoundTag.CODEC);
+
+    public static final StreamCodec<ByteBuf, AugmentSlot> AUGMENT_SLOT_STREAM_CODEC =
+            ByteBufCodecs.INT.map(MJRegistries.AUGMENT_SLOT::byId, MJRegistries.AUGMENT_SLOT::getId);
+}

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

@@ -0,0 +1,28 @@
+package com.portingdeadmods.modjam.utils;
+
+import com.google.common.collect.ImmutableMap;
+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.data.MJDataAttachments;
+import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
+import net.minecraft.world.entity.player.Player;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public final class AugmentHelper {
+    public static Augment getAugmentBySlot(Player player, AugmentSlot augmentSlot) {
+        return getAugments(player).get(augmentSlot);
+    }
+
+    public static Map<AugmentSlot, Augment> getAugments(Player player) {
+        return player.getData(MJDataAttachments.AUGMENTS);
+    }
+
+    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);
+    }
+}