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.ElytraModel;
import net.minecraft.client.model.EntityModel;
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.entity.RenderLayerParent;
import net.minecraft.client.renderer.entity.state.HumanoidRenderState;
import net.minecraft.client.renderer.entity.state.PlayerRenderState;
import net.minecraft.client.resources.PlayerSkin;
import net.minecraft.client.resources.model.EquipmentClientInfo;
import net.minecraft.core.component.DataComponents;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.equipment.EquipmentAsset;
import net.minecraft.world.item.equipment.Equippable;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class WingsLayer<S extends HumanoidRenderState, M extends EntityModel<S>> extends RenderLayer<S, M> {
	private final ElytraModel elytraModel;
	private final ElytraModel elytraBabyModel;
	private final EquipmentLayerRenderer equipmentRenderer;

	public WingsLayer(RenderLayerParent<S, M> renderer, EntityModelSet models, EquipmentLayerRenderer equipmentRenderer) {
		super(renderer);
		this.elytraModel = new ElytraModel(models.bakeLayer(ModelLayers.ELYTRA));
		this.elytraBabyModel = new ElytraModel(models.bakeLayer(ModelLayers.ELYTRA_BABY));
		this.equipmentRenderer = equipmentRenderer;
	}

	public void render(PoseStack poseStack, MultiBufferSource multiBufferSource, int i, S humanoidRenderState, float f, float g) {
		ItemStack itemStack = humanoidRenderState.chestEquipment;
		Equippable equippable = itemStack.get(DataComponents.EQUIPPABLE);
		if (equippable != null && !equippable.assetId().isEmpty()) {
			ResourceLocation resourceLocation = getPlayerElytraTexture(humanoidRenderState);
			ElytraModel elytraModel = humanoidRenderState.isBaby ? this.elytraBabyModel : this.elytraModel;
			poseStack.pushPose();
			poseStack.translate(0.0F, 0.0F, 0.125F);
			elytraModel.setupAnim(humanoidRenderState);
			this.equipmentRenderer
				.renderLayers(
					EquipmentClientInfo.LayerType.WINGS,
					(ResourceKey<EquipmentAsset>)equippable.assetId().get(),
					elytraModel,
					itemStack,
					poseStack,
					multiBufferSource,
					i,
					resourceLocation
				);
			poseStack.popPose();
		}
	}

	@Nullable
	private static ResourceLocation getPlayerElytraTexture(HumanoidRenderState renderState) {
		if (renderState instanceof PlayerRenderState playerRenderState) {
			PlayerSkin playerSkin = playerRenderState.skin;
			if (playerSkin.elytraTexture() != null) {
				return playerSkin.elytraTexture();
			}

			if (playerSkin.capeTexture() != null && playerRenderState.showCape) {
				return playerSkin.capeTexture();
			}
		}

		return null;
	}
}