1.21.3 WIP

6 files changed, 166 insertions, 96 deletions

diff --git a/src/main/kotlin/features/texturepack/BakedModelExtra.kt b/src/main/kotlin/features/texturepack/BakedModelExtra.kt
index 32f419a..6305748 100644
--- a/src/main/kotlin/features/texturepack/BakedModelExtra.kt
+++ b/src/main/kotlin/features/texturepack/BakedModelExtra.kt
@@ -1,11 +1,30 @@
-
 package moe.nea.firmament.features.texturepack
 
+import net.fabricmc.fabric.api.renderer.v1.model.WrapperBakedModel as WrapperBakedModelFabric
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.WrapperBakedModel
+import moe.nea.firmament.util.ErrorUtil
 
 interface BakedModelExtra {
+	companion object {
+		@JvmStatic
+		fun cast(originalModel: BakedModel): BakedModelExtra? {
+			var p = originalModel
+			for (i in 0..256) {
+				p = when (p) {
+					is BakedModelExtra -> return p
+					is WrapperBakedModel -> p.wrapped
+					is WrapperBakedModelFabric -> WrapperBakedModelFabric.unwrap(p)
+					else -> break
+				}
+			}
+			ErrorUtil.softError("Could not find a baked model for $originalModel")
+			return null
+		}
+	}
+
 	var tintOverrides_firmament: TintOverrides?
 
 	fun getHeadModel_firmament(): BakedModel?
-    fun setHeadModel_firmament(headModel: BakedModel?)
+	fun setHeadModel_firmament(headModel: BakedModel?)
 }
diff --git a/src/main/kotlin/features/texturepack/CustomBlockTextures.kt b/src/main/kotlin/features/texturepack/CustomBlockTextures.kt
index a149928..2f7f084 100644
--- a/src/main/kotlin/features/texturepack/CustomBlockTextures.kt
+++ b/src/main/kotlin/features/texturepack/CustomBlockTextures.kt
@@ -244,7 +244,7 @@ object CustomBlockTextures {
             .flatMap { it.lookup.values }
             .flatten()
             .mapTo(mutableSetOf()) { it.replacement.blockModelIdentifier }
-            .forEach { event.addNonItemModel(it) }
+            .forEach { event.addNonItemModel(it, it.id) }
     }
 
     private fun prepare(manager: ResourceManager): BakedReplacements {
@@ -263,7 +263,7 @@ object CustomBlockTextures {
                 val island = SkyBlockIsland.forMode(mode)
                 val islandMpa = map.getOrPut(island, ::mutableMapOf)
                 for ((blockId, replacement) in json.replacements) {
-                    val block = MC.defaultRegistries.getWrapperOrThrow(RegistryKeys.BLOCK)
+                    val block = MC.defaultRegistries.getOrThrow(RegistryKeys.BLOCK)
                         .getOptional(RegistryKey.of(RegistryKeys.BLOCK, blockId))
                         .getOrNull()
                     if (block == null) {
diff --git a/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt b/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
index 7b6e62b..54e9e11 100644
--- a/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
+++ b/src/main/kotlin/features/texturepack/CustomGlobalArmorOverrides.kt
@@ -3,13 +3,13 @@
 package moe.nea.firmament.features.texturepack
 
 import java.util.Optional
+import java.util.concurrent.atomic.AtomicInteger
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.Transient
 import kotlinx.serialization.UseSerializers
-import kotlin.jvm.optionals.getOrNull
-import net.minecraft.item.ArmorMaterial
 import net.minecraft.item.ItemStack
+import net.minecraft.item.equipment.EquipmentModel
 import net.minecraft.resource.ResourceManager
 import net.minecraft.resource.SinglePreparationResourceReloader
 import net.minecraft.util.Identifier
@@ -17,90 +17,139 @@ 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.collections.WeakCache
 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 = WeakCache.memoize<ItemStack, Optional<List<ArmorMaterial.Layer>>>("ArmorOverrides") { stack ->
-        val id = stack.skyBlockId ?: return@memoize Optional.empty()
-        val override = overrides[id.neuItem] ?: return@memoize Optional.empty()
-        for (suboverride in override.overrides) {
-            if (suboverride.predicate.test(stack)) {
-                return@memoize Optional.of(suboverride.bakedLayers)
-            }
-        }
-        return@memoize Optional.of(override.bakedLayers)
-    }
-
-    @JvmStatic
-    fun overrideArmor(stack: ItemStack): List<ArmorMaterial.Layer>? {
-        if (!CustomSkyBlockTextures.TConfig.enableArmorOverrides) return null
-        return overrideCache.invoke(stack).getOrNull()
-    }
-
-    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
-            }
-        })
-    }
+object CustomGlobalArmorOverrides {
+	@Serializable
+	data class ArmorOverride(
+		@SerialName("item_ids")
+		val itemIds: List<String>,
+		val layers: List<ArmorOverrideLayer>? = null,
+		val model: Identifier? = null,
+		val overrides: List<ArmorOverrideOverride> = listOf(),
+	) {
+		@Transient
+		lateinit var modelIdentifier: Identifier
+		fun bake() {
+			modelIdentifier = bakeModel(model, layers)
+			overrides.forEach { it.bake() }
+		}
+
+		init {
+			require(layers != null || model != null) { "Either model or layers must be specified for armor override" }
+			require(layers == null || model == null) { "Can't specify both model and layers for armor override" }
+		}
+	}
+
+	@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>? = null,
+		val model: Identifier? = null,
+	) {
+		init {
+			require(layers != null || model != null) { "Either model or layers must be specified for armor override override" }
+			require(layers == null || model == null) { "Can't specify both model and layers for armor override override" }
+		}
+
+		@Transient
+		lateinit var modelIdentifier: Identifier
+		fun bake() {
+			modelIdentifier = bakeModel(model, layers)
+		}
+	}
+
+
+	val overrideCache = WeakCache.memoize<ItemStack, Optional<Identifier>>("ArmorOverrides") { stack ->
+		val id = stack.skyBlockId ?: return@memoize Optional.empty()
+		val override = overrides[id.neuItem] ?: return@memoize Optional.empty()
+		for (suboverride in override.overrides) {
+			if (suboverride.predicate.test(stack)) {
+				return@memoize Optional.of(suboverride.modelIdentifier)
+			}
+		}
+		return@memoize Optional.of(override.modelIdentifier)
+	}
+
+	var overrides: Map<String, ArmorOverride> = mapOf()
+	private var bakedOverrides: MutableMap<Identifier, EquipmentModel> = mutableMapOf()
+	private val sentinelFirmRunning = AtomicInteger()
+
+	private fun bakeModel(model: Identifier?, layers: List<ArmorOverrideLayer>?): Identifier {
+		require(model == null || layers == null)
+		if (model != null) {
+			return model
+		} else if (layers != null) {
+			val idNumber = sentinelFirmRunning.incrementAndGet()
+			val identifier = Identifier.of("firmament:sentinel/$idNumber")
+			val equipmentLayers = layers.map {
+				EquipmentModel.Layer(
+					it.identifier, if (it.tint) {
+						Optional.of(EquipmentModel.Dyeable(Optional.empty()))
+					} else {
+						Optional.empty()
+					},
+					false
+				)
+			}
+			bakedOverrides[identifier] = EquipmentModel(
+				mapOf(
+					EquipmentModel.LayerType.HUMANOID to equipmentLayers,
+					EquipmentModel.LayerType.HUMANOID_LEGGINGS to equipmentLayers,
+				)
+			)
+			return identifier
+		} else {
+			error("Either model or layers must be non null")
+		}
+	}
+
+
+	@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) {
+				bakedOverrides.clear()
+				prepared.forEach { it.value.bake() }
+				overrides = prepared
+			}
+		})
+	}
+
+	@JvmStatic
+	fun overrideArmor(itemStack: ItemStack): Optional<Identifier> {
+		return overrideCache.invoke(itemStack)
+	}
+
+	@JvmStatic
+	fun overrideArmorLayer(id: Identifier): EquipmentModel? {
+		return bakedOverrides[id]
+	}
 
 }
diff --git a/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt b/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt
index 2771699..a1203df 100644
--- a/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt
+++ b/src/main/kotlin/features/texturepack/CustomGlobalTextures.kt
@@ -146,7 +146,7 @@ object CustomGlobalTextures : SinglePreparationResourceReloader<CustomGlobalText
                 it.overrides
                     .asSequence()
                     .filter { it.predicate.test(stack) }
-                    .map { models.modelManager.getModel(ModelIdentifier(it.model, "inventory")) }
+                    .map { models.getModel(it.model) }
                     .firstOrNull()
             }
             .intoOptional()
diff --git a/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt b/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt
index 0d0f8f2..9f641b8 100644
--- a/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt
+++ b/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt
@@ -1,9 +1,11 @@
 
 package moe.nea.firmament.features.texturepack
 
+import net.minecraft.client.render.model.Baker
 import net.minecraft.util.Identifier
 
 interface JsonUnbakedModelFirmExtra {
+	fun storeExtraBaker_firmament(baker: Baker)
 
     fun setHeadModel_firmament(identifier: Identifier?)
     fun getHeadModel_firmament(): Identifier?
diff --git a/src/main/kotlin/features/texturepack/TintOverrides.kt b/src/main/kotlin/features/texturepack/TintOverrides.kt
index 8006db8..85fcae4 100644
--- a/src/main/kotlin/features/texturepack/TintOverrides.kt
+++ b/src/main/kotlin/features/texturepack/TintOverrides.kt
@@ -3,7 +3,6 @@ package moe.nea.firmament.features.texturepack
 import com.google.gson.JsonObject
 import com.google.gson.JsonPrimitive
 import moe.nea.firmament.util.ErrorUtil
-import moe.nea.firmament.util.assertNotNullOr
 
 data class TintOverrides(
 	val layerMap: Map<Int, TintOverride> = mapOf()
@@ -14,21 +13,22 @@ data class TintOverrides(
 		val EMPTY = TintOverrides()
 		private val threadLocal = object : ThreadLocal<TintOverrides>() {}
 		fun enter(overrides: TintOverrides?) {
-			ErrorUtil.softCheck("Double entered tintOverrides") {
-				threadLocal.get() == null
-			}
+			ErrorUtil.softCheck("Double entered tintOverrides",
+			                    threadLocal.get() == null)
 			threadLocal.set(overrides ?: EMPTY)
 		}
 
 		fun exit(overrides: TintOverrides?) {
-			ErrorUtil.softCheck("Exited with non matching enter tintOverrides") {
-				threadLocal.get() == (overrides ?: EMPTY)
-			}
+			ErrorUtil.softCheck("Exited with non matching enter tintOverrides",
+			                    threadLocal.get() == (overrides ?: EMPTY))
 			threadLocal.remove()
 		}
 
-		fun getCurrentOverrides() =
-			assertNotNullOr(threadLocal.get(), "Got current tintOverrides without entering") { EMPTY }
+		fun getCurrentOverrides(): TintOverrides {
+			return ErrorUtil.notNullOr(threadLocal.get(), "Got current tintOverrides without entering") {
+				EMPTY
+			}
+		}
 
 		fun parse(jsonObject: JsonObject): TintOverrides {
 			val map = mutableMapOf<Int, TintOverride>()