package net.minecraft.client.resources.model;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import java.util.EnumMap;
import java.util.List;
import java.util.Map;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.block.model.BakedQuad;
import net.minecraft.client.renderer.block.model.BlockElement;
import net.minecraft.client.renderer.block.model.BlockElementFace;
import net.minecraft.client.renderer.block.model.FaceBakery;
import net.minecraft.client.renderer.block.model.ItemTransforms;
import net.minecraft.client.renderer.block.model.TextureSlots;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.core.Direction;
import net.minecraft.util.RandomSource;
import net.minecraft.world.level.block.state.BlockState;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class SimpleBakedModel implements BakedModel {
	public static final String PARTICLE_TEXTURE_REFERENCE = "particle";
	private final List<BakedQuad> unculledFaces;
	private final Map<Direction, List<BakedQuad>> culledFaces;
	private final boolean hasAmbientOcclusion;
	private final boolean isGui3d;
	private final boolean usesBlockLight;
	private final TextureAtlasSprite particleIcon;
	private final ItemTransforms transforms;

	public SimpleBakedModel(
		List<BakedQuad> unculledFaces,
		Map<Direction, List<BakedQuad>> culledFaces,
		boolean hasAmbientOcclusion,
		boolean useBlockLight,
		boolean isGui3d,
		TextureAtlasSprite particleIcon,
		ItemTransforms transforms
	) {
		this.unculledFaces = unculledFaces;
		this.culledFaces = culledFaces;
		this.hasAmbientOcclusion = hasAmbientOcclusion;
		this.isGui3d = isGui3d;
		this.usesBlockLight = useBlockLight;
		this.particleIcon = particleIcon;
		this.transforms = transforms;
	}

	public static BakedModel bakeElements(
		List<BlockElement> elements,
		TextureSlots textureSlots,
		SpriteGetter spriteGetter,
		ModelState modelState,
		boolean hasAmbientOcclusion,
		boolean useBlockLight,
		boolean isGui3d,
		ItemTransforms transforms
	) {
		TextureAtlasSprite textureAtlasSprite = findSprite(spriteGetter, textureSlots, "particle");
		SimpleBakedModel.Builder builder = new SimpleBakedModel.Builder(hasAmbientOcclusion, useBlockLight, isGui3d, transforms).particle(textureAtlasSprite);

		for (BlockElement blockElement : elements) {
			for (Direction direction : blockElement.faces.keySet()) {
				BlockElementFace blockElementFace = (BlockElementFace)blockElement.faces.get(direction);
				TextureAtlasSprite textureAtlasSprite2 = findSprite(spriteGetter, textureSlots, blockElementFace.texture());
				if (blockElementFace.cullForDirection() == null) {
					builder.addUnculledFace(bakeFace(blockElement, blockElementFace, textureAtlasSprite2, direction, modelState));
				} else {
					builder.addCulledFace(
						Direction.rotate(modelState.getRotation().getMatrix(), blockElementFace.cullForDirection()),
						bakeFace(blockElement, blockElementFace, textureAtlasSprite2, direction, modelState)
					);
				}
			}
		}

		return builder.build();
	}

	private static BakedQuad bakeFace(BlockElement element, BlockElementFace face, TextureAtlasSprite sprite, Direction facing, ModelState transform) {
		return FaceBakery.bakeQuad(element.from, element.to, face, sprite, facing, transform, element.rotation, element.shade, element.lightEmission);
	}

	private static TextureAtlasSprite findSprite(SpriteGetter spriteGetter, TextureSlots textureSlots, String material) {
		Material material2 = textureSlots.getMaterial(material);
		return material2 != null ? spriteGetter.get(material2) : spriteGetter.reportMissingReference(material);
	}

	@Override
	public List<BakedQuad> getQuads(@Nullable BlockState state, @Nullable Direction direction, RandomSource random) {
		return direction == null ? this.unculledFaces : (List)this.culledFaces.get(direction);
	}

	@Override
	public boolean useAmbientOcclusion() {
		return this.hasAmbientOcclusion;
	}

	@Override
	public boolean isGui3d() {
		return this.isGui3d;
	}

	@Override
	public boolean usesBlockLight() {
		return this.usesBlockLight;
	}

	@Override
	public TextureAtlasSprite getParticleIcon() {
		return this.particleIcon;
	}

	@Override
	public ItemTransforms getTransforms() {
		return this.transforms;
	}

	@Environment(EnvType.CLIENT)
	public static class Builder {
		private final ImmutableList.Builder<BakedQuad> unculledFaces = ImmutableList.builder();
		private final EnumMap<Direction, ImmutableList.Builder<BakedQuad>> culledFaces = Maps.newEnumMap(Direction.class);
		private final boolean hasAmbientOcclusion;
		@Nullable
		private TextureAtlasSprite particleIcon;
		private final boolean usesBlockLight;
		private final boolean isGui3d;
		private final ItemTransforms transforms;

		public Builder(boolean hasAmbientOcclusion, boolean useBlockLight, boolean isGui3d, ItemTransforms transforms) {
			this.hasAmbientOcclusion = hasAmbientOcclusion;
			this.usesBlockLight = useBlockLight;
			this.isGui3d = isGui3d;
			this.transforms = transforms;

			for (Direction direction : Direction.values()) {
				this.culledFaces.put(direction, ImmutableList.builder());
			}
		}

		public SimpleBakedModel.Builder addCulledFace(Direction facing, BakedQuad quad) {
			((ImmutableList.Builder)this.culledFaces.get(facing)).add(quad);
			return this;
		}

		public SimpleBakedModel.Builder addUnculledFace(BakedQuad quad) {
			this.unculledFaces.add(quad);
			return this;
		}

		public SimpleBakedModel.Builder particle(TextureAtlasSprite particleIcon) {
			this.particleIcon = particleIcon;
			return this;
		}

		public SimpleBakedModel.Builder item() {
			return this;
		}

		public BakedModel build() {
			if (this.particleIcon == null) {
				throw new RuntimeException("Missing particle!");
			} else {
				Map<Direction, List<BakedQuad>> map = Maps.transformValues(this.culledFaces, ImmutableList.Builder::build);
				return new SimpleBakedModel(
					this.unculledFaces.build(), new EnumMap(map), this.hasAmbientOcclusion, this.usesBlockLight, this.isGui3d, this.particleIcon, this.transforms
				);
			}
		}
	}
}