package net.minecraft.client.renderer.entity.layers;

import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.Util;
import net.minecraft.client.model.Model;
import net.minecraft.client.renderer.MultiBufferSource;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.client.renderer.Sheets;
import net.minecraft.client.renderer.entity.ItemRenderer;
import net.minecraft.client.renderer.texture.OverlayTexture;
import net.minecraft.client.renderer.texture.TextureAtlas;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.resources.model.EquipmentModelSet;
import net.minecraft.core.component.DataComponents;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.tags.ItemTags;
import net.minecraft.util.ARGB;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.component.DyedItemColor;
import net.minecraft.world.item.equipment.EquipmentModel;
import net.minecraft.world.item.equipment.trim.ArmorTrim;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class EquipmentLayerRenderer {
	private static final int NO_LAYER_COLOR = 0;
	private final EquipmentModelSet equipmentModels;
	private final Function<EquipmentLayerRenderer.LayerTextureKey, ResourceLocation> layerTextureLookup;
	private final Function<EquipmentLayerRenderer.TrimSpriteKey, TextureAtlasSprite> trimSpriteLookup;

	public EquipmentLayerRenderer(EquipmentModelSet equipmentModels, TextureAtlas trimsAtlas) {
		this.equipmentModels = equipmentModels;
		this.layerTextureLookup = Util.memoize(
			(Function<EquipmentLayerRenderer.LayerTextureKey, ResourceLocation>)(layerTextureKey -> layerTextureKey.layer.getTextureLocation(layerTextureKey.layerType))
		);
		this.trimSpriteLookup = Util.memoize((Function<EquipmentLayerRenderer.TrimSpriteKey, TextureAtlasSprite>)(trimSpriteKey -> {
			ResourceLocation resourceLocation = trimSpriteKey.trim.getTexture(trimSpriteKey.layerType, trimSpriteKey.equipmentModelId);
			return trimsAtlas.getSprite(resourceLocation);
		}));
	}

	public void renderLayers(
		EquipmentModel.LayerType layerType,
		ResourceLocation equipmentModel,
		Model armorModel,
		ItemStack item,
		PoseStack poseStack,
		MultiBufferSource bufferSource,
		int packedLight
	) {
		this.renderLayers(layerType, equipmentModel, armorModel, item, poseStack, bufferSource, packedLight, null);
	}

	public void renderLayers(
		EquipmentModel.LayerType layerType,
		ResourceLocation equipmentModel,
		Model armorModel,
		ItemStack item,
		PoseStack poseStack,
		MultiBufferSource bufferSource,
		int packedLight,
		@Nullable ResourceLocation playerTexture
	) {
		List<EquipmentModel.Layer> list = this.equipmentModels.get(equipmentModel).getLayers(layerType);
		if (!list.isEmpty()) {
			int i = item.is(ItemTags.DYEABLE) ? DyedItemColor.getOrDefault(item, 0) : 0;
			boolean bl = item.hasFoil();

			for (EquipmentModel.Layer layer : list) {
				int j = getColorForLayer(layer, i);
				if (j != 0) {
					ResourceLocation resourceLocation = layer.usePlayerTexture() && playerTexture != null
						? playerTexture
						: (ResourceLocation)this.layerTextureLookup.apply(new EquipmentLayerRenderer.LayerTextureKey(layerType, layer));
					VertexConsumer vertexConsumer = ItemRenderer.getArmorFoilBuffer(bufferSource, RenderType.armorCutoutNoCull(resourceLocation), bl);
					armorModel.renderToBuffer(poseStack, vertexConsumer, packedLight, OverlayTexture.NO_OVERLAY, j);
					bl = false;
				}
			}

			ArmorTrim armorTrim = item.get(DataComponents.TRIM);
			if (armorTrim != null) {
				TextureAtlasSprite textureAtlasSprite = (TextureAtlasSprite)this.trimSpriteLookup
					.apply(new EquipmentLayerRenderer.TrimSpriteKey(armorTrim, layerType, equipmentModel));
				VertexConsumer vertexConsumer2 = textureAtlasSprite.wrap(bufferSource.getBuffer(Sheets.armorTrimsSheet(armorTrim.pattern().value().decal())));
				armorModel.renderToBuffer(poseStack, vertexConsumer2, packedLight, OverlayTexture.NO_OVERLAY);
			}
		}
	}

	private static int getColorForLayer(EquipmentModel.Layer layer, int color) {
		Optional<EquipmentModel.Dyeable> optional = layer.dyeable();
		if (optional.isPresent()) {
			int i = (Integer)((EquipmentModel.Dyeable)optional.get()).colorWhenUndyed().map(ARGB::opaque).orElse(0);
			return color != 0 ? color : i;
		} else {
			return -1;
		}
	}

	@Environment(EnvType.CLIENT)
	record LayerTextureKey(EquipmentModel.LayerType layerType, EquipmentModel.Layer layer) {
	}

	@Environment(EnvType.CLIENT)
	record TrimSpriteKey(ArmorTrim trim, EquipmentModel.LayerType layerType, ResourceLocation equipmentModelId) {
	}
}