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

import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import java.util.Map;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.model.WolfModel;
import net.minecraft.client.model.geom.EntityModelSet;
import net.minecraft.client.model.geom.ModelLayers;
import net.minecraft.client.renderer.MultiBufferSource;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.client.renderer.entity.RenderLayerParent;
import net.minecraft.client.renderer.texture.OverlayTexture;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.tags.ItemTags;
import net.minecraft.util.FastColor;
import net.minecraft.world.entity.Crackiness;
import net.minecraft.world.entity.animal.Wolf;
import net.minecraft.world.item.AnimalArmorItem;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.component.DyedItemColor;

@Environment(EnvType.CLIENT)
public class WolfArmorLayer extends RenderLayer<Wolf, WolfModel<Wolf>> {
	private final WolfModel<Wolf> model;
	private static final Map<Crackiness.Level, ResourceLocation> ARMOR_CRACK_LOCATIONS = Map.of(
		Crackiness.Level.LOW,
		ResourceLocation.withDefaultNamespace("textures/entity/wolf/wolf_armor_crackiness_low.png"),
		Crackiness.Level.MEDIUM,
		ResourceLocation.withDefaultNamespace("textures/entity/wolf/wolf_armor_crackiness_medium.png"),
		Crackiness.Level.HIGH,
		ResourceLocation.withDefaultNamespace("textures/entity/wolf/wolf_armor_crackiness_high.png")
	);

	public WolfArmorLayer(RenderLayerParent<Wolf, WolfModel<Wolf>> renderer, EntityModelSet models) {
		super(renderer);
		this.model = new WolfModel<>(models.bakeLayer(ModelLayers.WOLF_ARMOR));
	}

	public void render(
		PoseStack poseStack,
		MultiBufferSource bufferSource,
		int packedLight,
		Wolf livingEntity,
		float limbSwing,
		float limbSwingAmount,
		float partialTick,
		float ageInTicks,
		float netHeadYaw,
		float headPitch
	) {
		if (livingEntity.hasArmor()) {
			ItemStack itemStack = livingEntity.getBodyArmorItem();
			if (itemStack.getItem() instanceof AnimalArmorItem animalArmorItem && animalArmorItem.getBodyType() == AnimalArmorItem.BodyType.CANINE) {
				this.getParentModel().copyPropertiesTo(this.model);
				this.model.prepareMobModel(livingEntity, limbSwing, limbSwingAmount, partialTick);
				this.model.setupAnim(livingEntity, limbSwing, limbSwingAmount, ageInTicks, netHeadYaw, headPitch);
				VertexConsumer vertexConsumer = bufferSource.getBuffer(RenderType.entityCutoutNoCull(animalArmorItem.getTexture()));
				this.model.renderToBuffer(poseStack, vertexConsumer, packedLight, OverlayTexture.NO_OVERLAY);
				this.maybeRenderColoredLayer(poseStack, bufferSource, packedLight, itemStack, animalArmorItem);
				this.maybeRenderCracks(poseStack, bufferSource, packedLight, itemStack);
			}
		}
	}

	private void maybeRenderColoredLayer(PoseStack poseStack, MultiBufferSource buffer, int packedLight, ItemStack armorStack, AnimalArmorItem armorItem) {
		if (armorStack.is(ItemTags.DYEABLE)) {
			int i = DyedItemColor.getOrDefault(armorStack, 0);
			if (FastColor.ARGB32.alpha(i) == 0) {
				return;
			}

			ResourceLocation resourceLocation = armorItem.getOverlayTexture();
			if (resourceLocation == null) {
				return;
			}

			this.model
				.renderToBuffer(
					poseStack, buffer.getBuffer(RenderType.entityCutoutNoCull(resourceLocation)), packedLight, OverlayTexture.NO_OVERLAY, FastColor.ARGB32.opaque(i)
				);
		}
	}

	private void maybeRenderCracks(PoseStack poseStack, MultiBufferSource buffer, int packedLight, ItemStack armorStack) {
		Crackiness.Level level = Crackiness.WOLF_ARMOR.byDamage(armorStack);
		if (level != Crackiness.Level.NONE) {
			ResourceLocation resourceLocation = (ResourceLocation)ARMOR_CRACK_LOCATIONS.get(level);
			VertexConsumer vertexConsumer = buffer.getBuffer(RenderType.entityTranslucent(resourceLocation));
			this.model.renderToBuffer(poseStack, vertexConsumer, packedLight, OverlayTexture.NO_OVERLAY);
		}
	}
}