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

import com.mojang.blaze3d.vertex.PoseStack;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.model.HumanoidModel;
import net.minecraft.client.renderer.MultiBufferSource;
import net.minecraft.client.renderer.entity.RenderLayerParent;
import net.minecraft.client.renderer.entity.state.HumanoidRenderState;
import net.minecraft.client.resources.model.EquipmentClientInfo;
import net.minecraft.core.component.DataComponents;
import net.minecraft.resources.ResourceKey;
import net.minecraft.world.entity.EquipmentSlot;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.equipment.EquipmentAsset;
import net.minecraft.world.item.equipment.Equippable;

@Environment(EnvType.CLIENT)
public class HumanoidArmorLayer<S extends HumanoidRenderState, M extends HumanoidModel<S>, A extends HumanoidModel<S>> extends RenderLayer<S, M> {
	private final A innerModel;
	private final A outerModel;
	private final A innerModelBaby;
	private final A outerModelBaby;
	private final EquipmentLayerRenderer equipmentRenderer;

	public HumanoidArmorLayer(RenderLayerParent<S, M> renderer, A innerModel, A outerModel, EquipmentLayerRenderer equipmentRenderer) {
		this(renderer, innerModel, outerModel, innerModel, outerModel, equipmentRenderer);
	}

	public HumanoidArmorLayer(
		RenderLayerParent<S, M> renderer, A innerModel, A outerModel, A innerModelBaby, A outerModelBaby, EquipmentLayerRenderer equipmentRenderer
	) {
		super(renderer);
		this.innerModel = innerModel;
		this.outerModel = outerModel;
		this.innerModelBaby = innerModelBaby;
		this.outerModelBaby = outerModelBaby;
		this.equipmentRenderer = equipmentRenderer;
	}

	public static boolean shouldRender(ItemStack stack, EquipmentSlot slot) {
		Equippable equippable = stack.get(DataComponents.EQUIPPABLE);
		return equippable != null && shouldRender(equippable, slot);
	}

	private static boolean shouldRender(Equippable equippable, EquipmentSlot slot) {
		return equippable.assetId().isPresent() && equippable.slot() == slot;
	}

	public void render(PoseStack poseStack, MultiBufferSource multiBufferSource, int i, S humanoidRenderState, float f, float g) {
		this.renderArmorPiece(
			poseStack, multiBufferSource, humanoidRenderState.chestEquipment, EquipmentSlot.CHEST, i, this.getArmorModel(humanoidRenderState, EquipmentSlot.CHEST)
		);
		this.renderArmorPiece(
			poseStack, multiBufferSource, humanoidRenderState.legsEquipment, EquipmentSlot.LEGS, i, this.getArmorModel(humanoidRenderState, EquipmentSlot.LEGS)
		);
		this.renderArmorPiece(
			poseStack, multiBufferSource, humanoidRenderState.feetEquipment, EquipmentSlot.FEET, i, this.getArmorModel(humanoidRenderState, EquipmentSlot.FEET)
		);
		this.renderArmorPiece(
			poseStack, multiBufferSource, humanoidRenderState.headEquipment, EquipmentSlot.HEAD, i, this.getArmorModel(humanoidRenderState, EquipmentSlot.HEAD)
		);
	}

	private void renderArmorPiece(PoseStack poseStack, MultiBufferSource bufferSource, ItemStack armorItem, EquipmentSlot slot, int packedLight, A model) {
		Equippable equippable = armorItem.get(DataComponents.EQUIPPABLE);
		if (equippable != null && shouldRender(equippable, slot)) {
			this.getParentModel().copyPropertiesTo(model);
			this.setPartVisibility(model, slot);
			EquipmentClientInfo.LayerType layerType = this.usesInnerModel(slot)
				? EquipmentClientInfo.LayerType.HUMANOID_LEGGINGS
				: EquipmentClientInfo.LayerType.HUMANOID;
			this.equipmentRenderer
				.renderLayers(layerType, (ResourceKey<EquipmentAsset>)equippable.assetId().orElseThrow(), model, armorItem, poseStack, bufferSource, packedLight);
		}
	}

	protected void setPartVisibility(A model, EquipmentSlot slot) {
		model.setAllVisible(false);
		switch (slot) {
			case HEAD:
				model.head.visible = true;
				model.hat.visible = true;
				break;
			case CHEST:
				model.body.visible = true;
				model.rightArm.visible = true;
				model.leftArm.visible = true;
				break;
			case LEGS:
				model.body.visible = true;
				model.rightLeg.visible = true;
				model.leftLeg.visible = true;
				break;
			case FEET:
				model.rightLeg.visible = true;
				model.leftLeg.visible = true;
		}
	}

	private A getArmorModel(S renderState, EquipmentSlot slot) {
		if (this.usesInnerModel(slot)) {
			return renderState.isBaby ? this.innerModelBaby : this.innerModel;
		} else {
			return renderState.isBaby ? this.outerModelBaby : this.outerModel;
		}
	}

	private boolean usesInnerModel(EquipmentSlot slot) {
		return slot == EquipmentSlot.LEGS;
	}
}