De-hardcode armor trim types

4 files changed, 95 insertions, 78 deletions

diff --git a/src/main/java/me/xmrvizzy/skyblocker/config/SkyblockerConfig.java b/src/main/java/me/xmrvizzy/skyblocker/config/SkyblockerConfig.java
index 056116a9..6c046d50 100644
--- a/src/main/java/me/xmrvizzy/skyblocker/config/SkyblockerConfig.java
+++ b/src/main/java/me/xmrvizzy/skyblocker/config/SkyblockerConfig.java
@@ -1,11 +1,10 @@
 package me.xmrvizzy.skyblocker.config;
 
-import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
-import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
-
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import me.shedaniel.autoconfig.AutoConfig;
 import me.shedaniel.autoconfig.ConfigData;
 import me.shedaniel.autoconfig.annotation.Config;
@@ -14,12 +13,14 @@ import me.shedaniel.autoconfig.serializer.ConfigSerializer;
 import me.shedaniel.autoconfig.serializer.GsonConfigSerializer;
 import me.xmrvizzy.skyblocker.SkyblockerMod;
 import me.xmrvizzy.skyblocker.chat.ChatFilterResult;
+import me.xmrvizzy.skyblocker.skyblock.item.CustomArmorTrims;
 import me.xmrvizzy.skyblocker.utils.Scheduler;
 import net.fabricmc.fabric.api.client.command.v2.ClientCommandRegistrationCallback;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
 import net.minecraft.client.resource.language.I18n;
 import net.minecraft.text.Style;
 import net.minecraft.text.Text;
+import net.minecraft.util.Identifier;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -196,12 +197,15 @@ public class SkyblockerConfig implements ConfigData {
 
         @ConfigEntry.Gui.Excluded
         public List<Integer> lockedSlots = new ArrayList<>();
-        
+
+        @ConfigEntry.Gui.Excluded
         public Object2ObjectOpenHashMap<String, Text> customItemNames = new Object2ObjectOpenHashMap<>();
-        
+
+        @ConfigEntry.Gui.Excluded
         public Object2IntOpenHashMap<String> customDyeColors = new Object2IntOpenHashMap<>();
-        
-        public Object2ObjectOpenHashMap<String, String> customArmorTrims = new Object2ObjectOpenHashMap<>();
+
+        @ConfigEntry.Gui.Excluded
+        public Object2ObjectOpenHashMap<String, CustomArmorTrims.ArmorTrimId> customArmorTrims = new Object2ObjectOpenHashMap<>();
     }
 
     public static class TabHudConf {
@@ -215,11 +219,10 @@ public class SkyblockerConfig implements ConfigData {
         @ConfigEntry.Gui.EnumHandler(option = ConfigEntry.Gui.EnumHandler.EnumDisplayOption.BUTTON)
         @ConfigEntry.Gui.Tooltip
         public NameSorting nameSorting = NameSorting.DEFAULT;
-        
     }
-    
+
     public enum NameSorting {
-    	DEFAULT,
+        DEFAULT,
         ALPHABETICAL;
 
         @Override
@@ -551,10 +554,11 @@ public class SkyblockerConfig implements ConfigData {
                 .setPrettyPrinting()
                 .registerTypeHierarchyAdapter(Text.class, new Text.Serializer())
                 .registerTypeHierarchyAdapter(Style.class, new Style.Serializer())
+                .registerTypeHierarchyAdapter(Identifier.class, new Identifier.Serializer())
                 .create();
-        
+
         ConfigSerializer.Factory<SkyblockerConfig> serializer = (cfg, cfgClass) -> new GsonConfigSerializer<>(cfg, cfgClass, gson);
-        
+
         AutoConfig.register(SkyblockerConfig.class, serializer);
         ClientCommandRegistrationCallback.EVENT.register(((dispatcher, registryAccess) -> dispatcher.register(literal(SkyblockerMod.NAMESPACE).then(optionsLiteral("config")).then(optionsLiteral("options")))));
     }
@@ -573,7 +577,7 @@ public class SkyblockerConfig implements ConfigData {
     public static SkyblockerConfig get() {
         return AutoConfig.getConfigHolder(SkyblockerConfig.class).getConfig();
     }
-    
+
     public static void save() {
         AutoConfig.getConfigHolder(SkyblockerConfig.class).save();
     }
diff --git a/src/main/java/me/xmrvizzy/skyblocker/mixin/ArmorTrimMixin.java b/src/main/java/me/xmrvizzy/skyblocker/mixin/ArmorTrimMixin.java
index bb8dde66..a90cf4d4 100644
--- a/src/main/java/me/xmrvizzy/skyblocker/mixin/ArmorTrimMixin.java
+++ b/src/main/java/me/xmrvizzy/skyblocker/mixin/ArmorTrimMixin.java
@@ -1,13 +1,7 @@
 package me.xmrvizzy.skyblocker.mixin;
 
-import java.util.Optional;
-
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.injection.At;
-
 import com.llamalad7.mixinextras.injector.ModifyReturnValue;
 import com.llamalad7.mixinextras.sugar.Local;
-
 import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import me.xmrvizzy.skyblocker.config.SkyblockerConfig;
 import me.xmrvizzy.skyblocker.skyblock.item.CustomArmorTrims;
@@ -15,25 +9,29 @@ import me.xmrvizzy.skyblocker.utils.Utils;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.trim.ArmorTrim;
 import net.minecraft.nbt.NbtCompound;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+
+import java.util.Optional;
 
 @Mixin(ArmorTrim.class)
 public class ArmorTrimMixin {
 
 	@ModifyReturnValue(method = "getTrim", at = @At("RETURN"))
-	private static Optional<ArmorTrim> skyblocker$customArmorTrims(Optional<ArmorTrim> original, @Local ItemStack stack) {
+	private static Optional<ArmorTrim> skyblocker$customArmorTrims(@SuppressWarnings("OptionalUsedAsFieldOrParameterType") Optional<ArmorTrim> original, @Local ItemStack stack) {
 		NbtCompound nbt = stack.getNbt();
 
 		if (Utils.isOnSkyblock() && nbt != null && nbt.contains("ExtraAttributes")) {
-			Object2ObjectOpenHashMap<String, String> customTrims = SkyblockerConfig.get().general.customArmorTrims;
+			Object2ObjectOpenHashMap<String, CustomArmorTrims.ArmorTrimId> customTrims = SkyblockerConfig.get().general.customArmorTrims;
 			NbtCompound extraAttributes = nbt.getCompound("ExtraAttributes");
 			String itemUuid = extraAttributes.contains("uuid") ? extraAttributes.getString("uuid") : null;
-			
-			if(customTrims.containsKey(itemUuid)) {
-				String trimKey = customTrims.get(itemUuid);
+
+			if (customTrims.containsKey(itemUuid)) {
+				CustomArmorTrims.ArmorTrimId trimKey = customTrims.get(itemUuid);
 				return CustomArmorTrims.TRIMS_CACHE.getOrDefault(trimKey, original);
 			}
 		}
-		
+
 		return original;
 	}
 }
diff --git a/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorDyeColors.java b/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorDyeColors.java
index 46552f65..7c4ed423 100644
--- a/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorDyeColors.java
+++ b/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorDyeColors.java
@@ -10,7 +10,7 @@ import net.fabricmc.fabric.api.client.command.v2.ClientCommandManager;
 import net.fabricmc.fabric.api.client.command.v2.ClientCommandRegistrationCallback;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
 import net.minecraft.command.CommandRegistryAccess;
-import net.minecraft.item.DyeableArmorItem;
+import net.minecraft.item.DyeableItem;
 import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.text.Text;
@@ -40,7 +40,7 @@ public class CustomArmorDyeColors {
 		}
 
 		if (Utils.isOnSkyblock() && heldItem != null) {
-			if (heldItem.getItem() instanceof DyeableArmorItem) {
+			if (heldItem.getItem() instanceof DyeableItem) {
 				if (nbt != null && nbt.contains("ExtraAttributes")) {
 					NbtCompound extraAttributes = nbt.getCompound("ExtraAttributes");
 					String itemUuid = extraAttributes.contains("uuid") ? extraAttributes.getString("uuid") : null;
diff --git a/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorTrims.java b/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorTrims.java
index 3adfe66e..ffc3f67d 100644
--- a/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorTrims.java
+++ b/src/main/java/me/xmrvizzy/skyblocker/skyblock/item/CustomArmorTrims.java
@@ -1,15 +1,9 @@
 package me.xmrvizzy.skyblocker.skyblock.item;
 
-import java.util.Arrays;
-import java.util.Optional;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
-import com.mojang.brigadier.arguments.StringArgumentType;
-
+import com.mojang.brigadier.suggestion.SuggestionProvider;
+import it.unimi.dsi.fastutil.Pair;
 import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import me.xmrvizzy.skyblocker.config.SkyblockerConfig;
 import me.xmrvizzy.skyblocker.utils.SkyblockEvents;
@@ -18,74 +12,82 @@ import net.fabricmc.fabric.api.client.command.v2.ClientCommandManager;
 import net.fabricmc.fabric.api.client.command.v2.ClientCommandRegistrationCallback;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
 import net.minecraft.client.MinecraftClient;
+import net.minecraft.client.network.ClientPlayerEntity;
 import net.minecraft.command.CommandRegistryAccess;
 import net.minecraft.command.CommandSource;
+import net.minecraft.command.argument.IdentifierArgumentType;
 import net.minecraft.item.ArmorItem;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.trim.ArmorTrim;
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.nbt.NbtOps;
-import net.minecraft.nbt.NbtString;
-import net.minecraft.registry.DynamicRegistryManager;
-import net.minecraft.registry.RegistryOps;
+import net.minecraft.registry.*;
 import net.minecraft.text.Text;
+import net.minecraft.util.Identifier;
+import org.jetbrains.annotations.NotNull;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Optional;
 
 public class CustomArmorTrims {
 	private static final Logger LOGGER = LoggerFactory.getLogger(CustomArmorTrims.class);
-	public static final Object2ObjectOpenHashMap<String, Optional<ArmorTrim>> TRIMS_CACHE = new Object2ObjectOpenHashMap<>();
-	private static final String[] TRIM_MATERIALS = { "quartz", "iron", "netherite", "redstone", "copper", "gold", "emerald", "diamond", "lapis", "amethyst" };
-	private static final String[] TRIM_PATTERNS = { "sentry", "dune", "coast", "wild", "ward", "eye", "vex", "tide", "snout", "rib", "spire", "wayfinder", "shaper", "silence", "raiser", "host" };
-	
+	public static final Object2ObjectOpenHashMap<ArmorTrimId, Optional<ArmorTrim>> TRIMS_CACHE = new Object2ObjectOpenHashMap<>();
 	private static boolean trimsInitialized = false;
-	
+
 	public static void init() {
 		SkyblockEvents.JOIN.register(CustomArmorTrims::initializeTrimCache);
 		ClientCommandRegistrationCallback.EVENT.register(CustomArmorTrims::registerCommand);
 	}
-	
+
 	private static void initializeTrimCache() {
-		if (!trimsInitialized) {
-			MinecraftClient client = MinecraftClient.getInstance();
-			DynamicRegistryManager registryManager = client.player.networkHandler.getRegistryManager();
-			
-			for (String material: TRIM_MATERIALS) {
-				for (String pattern: TRIM_PATTERNS) {
-					String key = material + "+" + pattern;
-					
+		ClientPlayerEntity player = MinecraftClient.getInstance().player;
+		if (!trimsInitialized && player != null) {
+			TRIMS_CACHE.clear();
+			DynamicRegistryManager registryManager = player.networkHandler.getRegistryManager();
+			for (Identifier material : registryManager.get(RegistryKeys.TRIM_MATERIAL).getIds()) {
+				for (Identifier pattern : registryManager.get(RegistryKeys.TRIM_PATTERN).getIds()) {
 					NbtCompound compound = new NbtCompound();
-					compound.put("material", NbtString.of("minecraft:" + material));
-					compound.put("pattern", NbtString.of("minecraft:" + pattern));
-					
+					compound.putString("material", material.toString());
+					compound.putString("pattern", pattern.toString());
+
 					ArmorTrim trim = ArmorTrim.CODEC.parse(RegistryOps.of(NbtOps.INSTANCE, registryManager), compound).resultOrPartial(LOGGER::error).orElse(null);
-					
-					//Something went terribly wrong
+
+					// Something went terribly wrong
 					if (trim == null) throw new IllegalStateException("Trim shouldn't be null! [" + "\"" + material + "\",\"" + pattern + "\"]");
-					
-					TRIMS_CACHE.put(key, Optional.of(trim));
+
+					TRIMS_CACHE.put(new ArmorTrimId(material, pattern), Optional.of(trim));
 				}
 			}
-			
+
 			LOGGER.info("[Skyblocker] Successfully cached all armor trims!");
 			trimsInitialized = true;
 		}
 	}
-	
+
 	private static void registerCommand(CommandDispatcher<FabricClientCommandSource> dispatcher, CommandRegistryAccess registryAccess) {
 		dispatcher.register(ClientCommandManager.literal("skyblocker")
 				.then(ClientCommandManager.literal("custom")
 						.then(ClientCommandManager.literal("armorTrim")
 								.executes(context -> customizeTrim(context.getSource(), null, null))
-								.then(ClientCommandManager.argument("material", StringArgumentType.word())
-										.suggests((context, builder) -> CommandSource.suggestMatching(TRIM_MATERIALS, builder))
-										.then(ClientCommandManager.argument("pattern", StringArgumentType.word())
-												.suggests((context, builder) -> CommandSource.suggestMatching(TRIM_PATTERNS, builder))
-												.executes(context -> customizeTrim(context.getSource(), StringArgumentType.getString(context, "material").toLowerCase(), StringArgumentType.getString(context, "pattern").toLowerCase())))))));
+								.then(ClientCommandManager.argument("material", IdentifierArgumentType.identifier())
+										.suggests(getIdSuggestionProvider(RegistryKeys.TRIM_MATERIAL))
+										.executes(context -> customizeTrim(context.getSource(), context.getArgument("material", Identifier.class), null))
+										.then(ClientCommandManager.argument("pattern", IdentifierArgumentType.identifier())
+												.suggests(getIdSuggestionProvider(RegistryKeys.TRIM_PATTERN))
+												.executes(context -> customizeTrim(context.getSource(), context.getArgument("material", Identifier.class), context.getArgument("pattern", Identifier.class))))))));
+	}
+
+	@NotNull
+	private static SuggestionProvider<FabricClientCommandSource> getIdSuggestionProvider(RegistryKey<? extends Registry<?>> registryKey) {
+		return (context, builder) -> context.getSource().listIdSuggestions(registryKey, CommandSource.SuggestedIdType.ELEMENTS, builder, context);
 	}
-	
-	private static int customizeTrim(FabricClientCommandSource source, String material, String pattern) {
+
+	@SuppressWarnings("SameReturnValue")
+	private static int customizeTrim(FabricClientCommandSource source, Identifier material, Identifier pattern) {
 		ItemStack heldItem = source.getPlayer().getMainHandStack();
 		NbtCompound nbt = (heldItem != null) ? heldItem.getNbt() : null;
-		
+
 		if (Utils.isOnSkyblock() && heldItem != null) {
 			if (heldItem.getItem() instanceof ArmorItem) {
 				if (nbt != null && nbt.contains("ExtraAttributes")) {
@@ -93,8 +95,8 @@ public class CustomArmorTrims {
 					String itemUuid = extraAttributes.contains("uuid") ? extraAttributes.getString("uuid") : null;
 
 					if (itemUuid != null) {
-						Object2ObjectOpenHashMap<String, String> customArmorTrims = SkyblockerConfig.get().general.customArmorTrims;
-						
+						Object2ObjectOpenHashMap<String, ArmorTrimId> customArmorTrims = SkyblockerConfig.get().general.customArmorTrims;
+
 						if (material == null && pattern == null) {
 							if (customArmorTrims.containsKey(itemUuid)) {
 								customArmorTrims.remove(itemUuid);
@@ -104,15 +106,16 @@ public class CustomArmorTrims {
 								source.sendFeedback(Text.translatable("skyblocker.customArmorTrims.neverHad"));
 							}
 						} else {
-							
-							//Ensure that the material & trim are valid
-							if (!Arrays.stream(TRIM_MATERIALS).anyMatch(material::equals) || !Arrays.stream(TRIM_PATTERNS).anyMatch(pattern::equals)) {
+							// Ensure that the material & trim are valid
+							ArmorTrimId trimId = new ArmorTrimId(material, pattern);
+							if (TRIMS_CACHE.get(trimId) == null) {
 								source.sendError(Text.translatable("skyblocker.customArmorTrims.invalidMaterialOrPattern"));
-								
+
 								return Command.SINGLE_SUCCESS;
 							}
-							
-							customArmorTrims.put(itemUuid, material + "+" + pattern);
+
+							customArmorTrims.put(itemUuid, trimId);
+							SkyblockerConfig.save();
 							source.sendFeedback(Text.translatable("skyblocker.customArmorTrims.added"));
 						}
 					} else {
@@ -126,7 +129,19 @@ public class CustomArmorTrims {
 		} else {
 			source.sendError(Text.translatable("skyblocker.customArmorTrims.unableToSetTrim"));
 		}
-		
+
 		return Command.SINGLE_SUCCESS;
 	}
+
+	public record ArmorTrimId(Identifier material, Identifier pattern) implements Pair<Identifier, Identifier> {
+		@Override
+		public Identifier left() {
+			return material();
+		}
+
+		@Override
+		public Identifier right() {
+			return pattern();
+		}
+	}
 }