package net.minecraft.client.renderer.special;

import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import com.mojang.serialization.MapCodec;
import java.util.Objects;
import java.util.Set;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.model.ShieldModel;
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.blockentity.BannerRenderer;
import net.minecraft.client.renderer.entity.ItemRenderer;
import net.minecraft.client.resources.model.Material;
import net.minecraft.client.resources.model.ModelBakery;
import net.minecraft.core.component.DataComponentMap;
import net.minecraft.core.component.DataComponents;
import net.minecraft.world.item.DyeColor;
import net.minecraft.world.item.ItemDisplayContext;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.block.entity.BannerPatternLayers;
import org.jetbrains.annotations.Nullable;
import org.joml.Vector3f;

@Environment(EnvType.CLIENT)
public class ShieldSpecialRenderer implements SpecialModelRenderer<DataComponentMap> {
	private final ShieldModel model;

	public ShieldSpecialRenderer(ShieldModel model) {
		this.model = model;
	}

	@Nullable
	public DataComponentMap extractArgument(ItemStack stack) {
		return stack.immutableComponents();
	}

	public void render(
		@Nullable DataComponentMap patterns,
		ItemDisplayContext displayContext,
		PoseStack poseStack,
		MultiBufferSource bufferSource,
		int packedLight,
		int packedOverlay,
		boolean hasFoilType
	) {
		BannerPatternLayers bannerPatternLayers = patterns != null
			? patterns.getOrDefault(DataComponents.BANNER_PATTERNS, BannerPatternLayers.EMPTY)
			: BannerPatternLayers.EMPTY;
		DyeColor dyeColor = patterns != null ? patterns.get(DataComponents.BASE_COLOR) : null;
		boolean bl = !bannerPatternLayers.layers().isEmpty() || dyeColor != null;
		poseStack.pushPose();
		poseStack.scale(1.0F, -1.0F, -1.0F);
		Material material = bl ? ModelBakery.SHIELD_BASE : ModelBakery.NO_PATTERN_SHIELD;
		VertexConsumer vertexConsumer = material.sprite()
			.wrap(ItemRenderer.getFoilBuffer(bufferSource, this.model.renderType(material.atlasLocation()), displayContext == ItemDisplayContext.GUI, hasFoilType));
		this.model.handle().render(poseStack, vertexConsumer, packedLight, packedOverlay);
		if (bl) {
			BannerRenderer.renderPatterns(
				poseStack,
				bufferSource,
				packedLight,
				packedOverlay,
				this.model.plate(),
				material,
				false,
				(DyeColor)Objects.requireNonNullElse(dyeColor, DyeColor.WHITE),
				bannerPatternLayers,
				hasFoilType,
				false
			);
		} else {
			this.model.plate().render(poseStack, vertexConsumer, packedLight, packedOverlay);
		}

		poseStack.popPose();
	}

	@Override
	public void getExtents(Set<Vector3f> output) {
		PoseStack poseStack = new PoseStack();
		poseStack.scale(1.0F, -1.0F, -1.0F);
		this.model.root().getExtentsForGui(poseStack, output);
	}

	@Environment(EnvType.CLIENT)
	public record Unbaked() implements SpecialModelRenderer.Unbaked {
		public static final ShieldSpecialRenderer.Unbaked INSTANCE = new ShieldSpecialRenderer.Unbaked();
		public static final MapCodec<ShieldSpecialRenderer.Unbaked> MAP_CODEC = MapCodec.unit(INSTANCE);

		@Override
		public MapCodec<ShieldSpecialRenderer.Unbaked> type() {
			return MAP_CODEC;
		}

		@Override
		public SpecialModelRenderer<?> bake(EntityModelSet modelSet) {
			return new ShieldSpecialRenderer(new ShieldModel(modelSet.bakeLayer(ModelLayers.SHIELD)));
		}
	}
}