De-hardcode armor trim types

2 files changed, 68 insertions, 53 deletions

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();
+		}
+	}
 }