package net.minecraft.client.renderer.item;

import com.mojang.blaze3d.vertex.PoseStack;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Supplier;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.MultiBufferSource;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.client.renderer.block.model.BakedQuad;
import net.minecraft.client.renderer.block.model.ItemTransform;
import net.minecraft.client.renderer.entity.ItemRenderer;
import net.minecraft.client.renderer.special.SpecialModelRenderer;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.util.RandomSource;
import net.minecraft.world.item.ItemDisplayContext;
import org.jetbrains.annotations.Nullable;
import org.joml.Matrix4f;
import org.joml.Vector3f;
import org.joml.Vector3fc;

@Environment(EnvType.CLIENT)
public class ItemStackRenderState {
	ItemDisplayContext displayContext = ItemDisplayContext.NONE;
	private int activeLayerCount;
	private ItemStackRenderState.LayerRenderState[] layers = new ItemStackRenderState.LayerRenderState[]{new ItemStackRenderState.LayerRenderState()};

	public void ensureCapacity(int expectedSize) {
		int i = this.layers.length;
		int j = this.activeLayerCount + expectedSize;
		if (j > i) {
			this.layers = (ItemStackRenderState.LayerRenderState[])Arrays.copyOf(this.layers, j);

			for (int k = i; k < j; k++) {
				this.layers[k] = new ItemStackRenderState.LayerRenderState();
			}
		}
	}

	public ItemStackRenderState.LayerRenderState newLayer() {
		this.ensureCapacity(1);
		return this.layers[this.activeLayerCount++];
	}

	public void clear() {
		this.displayContext = ItemDisplayContext.NONE;

		for (int i = 0; i < this.activeLayerCount; i++) {
			this.layers[i].clear();
		}

		this.activeLayerCount = 0;
	}

	private ItemStackRenderState.LayerRenderState firstLayer() {
		return this.layers[0];
	}

	public boolean isEmpty() {
		return this.activeLayerCount == 0;
	}

	public boolean usesBlockLight() {
		return this.firstLayer().usesBlockLight;
	}

	@Nullable
	public TextureAtlasSprite pickParticleIcon(RandomSource random) {
		return this.activeLayerCount == 0 ? null : this.layers[random.nextInt(this.activeLayerCount)].particleIcon;
	}

	public void visitExtents(Consumer<Vector3fc> visitor) {
		Vector3f vector3f = new Vector3f();
		PoseStack.Pose pose = new PoseStack.Pose();

		for (int i = 0; i < this.activeLayerCount; i++) {
			ItemStackRenderState.LayerRenderState layerRenderState = this.layers[i];
			layerRenderState.transform.apply(this.displayContext.leftHand(), pose);
			Matrix4f matrix4f = pose.pose();
			Vector3f[] vector3fs = (Vector3f[])layerRenderState.extents.get();

			for (Vector3f vector3f2 : vector3fs) {
				visitor.accept(vector3f.set(vector3f2).mulPosition(matrix4f));
			}

			pose.setIdentity();
		}
	}

	public void render(PoseStack poseStack, MultiBufferSource bufferSource, int packedLight, int packedOverlay) {
		for (int i = 0; i < this.activeLayerCount; i++) {
			this.layers[i].render(poseStack, bufferSource, packedLight, packedOverlay);
		}
	}

	@Environment(EnvType.CLIENT)
	public static enum FoilType {
		NONE,
		STANDARD,
		SPECIAL;
	}

	@Environment(EnvType.CLIENT)
	public class LayerRenderState {
		private static final Vector3f[] NO_EXTENTS = new Vector3f[0];
		public static final Supplier<Vector3f[]> NO_EXTENTS_SUPPLIER = () -> NO_EXTENTS;
		private final List<BakedQuad> quads = new ArrayList();
		boolean usesBlockLight;
		@Nullable
		TextureAtlasSprite particleIcon;
		ItemTransform transform = ItemTransform.NO_TRANSFORM;
		@Nullable
		private RenderType renderType;
		private ItemStackRenderState.FoilType foilType = ItemStackRenderState.FoilType.NONE;
		private int[] tintLayers = new int[0];
		@Nullable
		private SpecialModelRenderer<Object> specialRenderer;
		@Nullable
		private Object argumentForSpecialRendering;
		Supplier<Vector3f[]> extents = NO_EXTENTS_SUPPLIER;

		public void clear() {
			this.quads.clear();
			this.renderType = null;
			this.foilType = ItemStackRenderState.FoilType.NONE;
			this.specialRenderer = null;
			this.argumentForSpecialRendering = null;
			Arrays.fill(this.tintLayers, -1);
			this.usesBlockLight = false;
			this.particleIcon = null;
			this.transform = ItemTransform.NO_TRANSFORM;
			this.extents = NO_EXTENTS_SUPPLIER;
		}

		public List<BakedQuad> prepareQuadList() {
			return this.quads;
		}

		public void setRenderType(RenderType renderType) {
			this.renderType = renderType;
		}

		public void setUsesBlockLight(boolean usesBlockLight) {
			this.usesBlockLight = usesBlockLight;
		}

		public void setExtents(Supplier<Vector3f[]> extents) {
			this.extents = extents;
		}

		public void setParticleIcon(TextureAtlasSprite particleIcon) {
			this.particleIcon = particleIcon;
		}

		public void setTransform(ItemTransform transform) {
			this.transform = transform;
		}

		public <T> void setupSpecialModel(SpecialModelRenderer<T> renderer, @Nullable T argument) {
			this.specialRenderer = eraseSpecialRenderer(renderer);
			this.argumentForSpecialRendering = argument;
		}

		private static SpecialModelRenderer<Object> eraseSpecialRenderer(SpecialModelRenderer<?> specialRenderer) {
			return (SpecialModelRenderer<Object>)specialRenderer;
		}

		public void setFoilType(ItemStackRenderState.FoilType foilType) {
			this.foilType = foilType;
		}

		public int[] prepareTintLayers(int count) {
			if (count > this.tintLayers.length) {
				this.tintLayers = new int[count];
				Arrays.fill(this.tintLayers, -1);
			}

			return this.tintLayers;
		}

		void render(PoseStack poseStack, MultiBufferSource bufferSource, int packedLight, int packedOverlay) {
			poseStack.pushPose();
			this.transform.apply(ItemStackRenderState.this.displayContext.leftHand(), poseStack.last());
			if (this.specialRenderer != null) {
				this.specialRenderer
					.render(
						this.argumentForSpecialRendering,
						ItemStackRenderState.this.displayContext,
						poseStack,
						bufferSource,
						packedLight,
						packedOverlay,
						this.foilType != ItemStackRenderState.FoilType.NONE
					);
			} else if (this.renderType != null) {
				ItemRenderer.renderItem(
					ItemStackRenderState.this.displayContext, poseStack, bufferSource, packedLight, packedOverlay, this.tintLayers, this.quads, this.renderType, this.foilType
				);
			}

			poseStack.popPose();
		}
	}
}