1 files changed, 106 insertions, 0 deletions
diff --git a/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt b/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
new file mode 100644
index 0000000..23577ee
--- /dev/null
+++ b/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
@@ -0,0 +1,106 @@
+
+@file:UseSerializers(IdentifierSerializer::class)
+
+package moe.nea.firmament.features.texturepack
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.Transient
+import kotlinx.serialization.UseSerializers
+import net.minecraft.item.ArmorMaterial
+import net.minecraft.item.ItemStack
+import net.minecraft.resource.ResourceManager
+import net.minecraft.resource.SinglePreparationResourceReloader
+import net.minecraft.util.Identifier
+import net.minecraft.util.profiler.Profiler
+import moe.nea.firmament.Firmament
+import moe.nea.firmament.annotations.Subscribe
+import moe.nea.firmament.events.FinalizeResourceManagerEvent
+import moe.nea.firmament.events.subscription.SubscriptionOwner
+import moe.nea.firmament.features.FirmamentFeature
+import moe.nea.firmament.features.texturepack.CustomGlobalTextures.logger
+import moe.nea.firmament.util.IdentifierSerializer
+import moe.nea.firmament.util.IdentityCharacteristics
+import moe.nea.firmament.util.computeNullableFunction
+import moe.nea.firmament.util.skyBlockId
+
+object CustomGlobalArmorOverrides : SubscriptionOwner {
+    @Serializable
+    data class ArmorOverride(
+        @SerialName("item_ids")
+        val itemIds: List<String>,
+        val layers: List<ArmorOverrideLayer>,
+        val overrides: List<ArmorOverrideOverride> = listOf(),
+    ) {
+        @Transient
+        val bakedLayers = bakeLayers(layers)
+    }
+
+    fun bakeLayers(layers: List<ArmorOverrideLayer>): List<ArmorMaterial.Layer> {
+        return layers.map { ArmorMaterial.Layer(it.identifier, it.suffix, it.tint) }
+    }
+
+    @Serializable
+    data class ArmorOverrideLayer(
+        val tint: Boolean = false,
+        val identifier: Identifier,
+        val suffix: String = "",
+    )
+
+    @Serializable
+    data class ArmorOverrideOverride(
+        val predicate: FirmamentModelPredicate,
+        val layers: List<ArmorOverrideLayer>,
+    ) {
+        @Transient
+        val bakedLayers = bakeLayers(layers)
+    }
+
+    override val delegateFeature: FirmamentFeature
+        get() = CustomSkyBlockTextures
+
+    val overrideCache = mutableMapOf<IdentityCharacteristics<ItemStack>, Any>()
+
+    @JvmStatic
+    fun overrideArmor(stack: ItemStack): List<ArmorMaterial.Layer>? {
+        if (!CustomSkyBlockTextures.TConfig.enableArmorOverrides) return null
+        return overrideCache.computeNullableFunction(IdentityCharacteristics(stack)) {
+            val id = stack.skyBlockId ?: return@computeNullableFunction null
+            val override = overrides[id.neuItem] ?: return@computeNullableFunction null
+            for (suboverride in override.overrides) {
+                if (suboverride.predicate.test(stack)) {
+                    return@computeNullableFunction suboverride.bakedLayers
+                }
+            }
+            return@computeNullableFunction override.bakedLayers
+        }
+    }
+
+    var overrides: Map<String, ArmorOverride> = mapOf()
+
+    @Subscribe
+    fun onStart(event: FinalizeResourceManagerEvent) {
+        event.resourceManager.registerReloader(object :
+                                                   SinglePreparationResourceReloader<Map<String, ArmorOverride>>() {
+            override fun prepare(manager: ResourceManager, profiler: Profiler): Map<String, ArmorOverride> {
+                val overrideFiles = manager.findResources("overrides/armor_models") {
+                    it.namespace == "firmskyblock" && it.path.endsWith(".json")
+                }
+                val overrides = overrideFiles.mapNotNull {
+                    Firmament.tryDecodeJsonFromStream<ArmorOverride>(it.value.inputStream).getOrElse { ex ->
+                        logger.error("Failed to load armor texture override at ${it.key}", ex)
+                        null
+                    }
+                }
+                val associatedMap = overrides.flatMap { obj -> obj.itemIds.map { it to obj } }
+                    .toMap()
+                return associatedMap
+            }
+
+            override fun apply(prepared: Map<String, ArmorOverride>, manager: ResourceManager, profiler: Profiler) {
+                overrides = prepared
+            }
+        })
+    }
+
+}