package net.minecraft.client.resources.model;

import com.mojang.logging.LogUtils;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.function.Function;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.model.geom.EntityModelSet;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.client.renderer.Sheets;
import net.minecraft.client.renderer.block.model.BlockStateModel;
import net.minecraft.client.renderer.block.model.ItemTransforms;
import net.minecraft.client.renderer.block.model.SimpleModelWrapper;
import net.minecraft.client.renderer.block.model.SingleVariant;
import net.minecraft.client.renderer.block.model.TextureSlots;
import net.minecraft.client.renderer.item.ClientItem;
import net.minecraft.client.renderer.item.ItemModel;
import net.minecraft.client.renderer.item.MissingItemModel;
import net.minecraft.client.renderer.item.ModelRenderProperties;
import net.minecraft.client.renderer.texture.TextureAtlas;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.resources.model.ModelBakery.MissingModels.1;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.util.thread.ParallelMapTransform;
import net.minecraft.world.level.block.state.BlockState;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class ModelBakery {
	public static final Material FIRE_0 = new Material(TextureAtlas.LOCATION_BLOCKS, ResourceLocation.withDefaultNamespace("block/fire_0"));
	public static final Material FIRE_1 = new Material(TextureAtlas.LOCATION_BLOCKS, ResourceLocation.withDefaultNamespace("block/fire_1"));
	public static final Material LAVA_FLOW = new Material(TextureAtlas.LOCATION_BLOCKS, ResourceLocation.withDefaultNamespace("block/lava_flow"));
	public static final Material WATER_FLOW = new Material(TextureAtlas.LOCATION_BLOCKS, ResourceLocation.withDefaultNamespace("block/water_flow"));
	public static final Material WATER_OVERLAY = new Material(TextureAtlas.LOCATION_BLOCKS, ResourceLocation.withDefaultNamespace("block/water_overlay"));
	public static final Material BANNER_BASE = new Material(Sheets.BANNER_SHEET, ResourceLocation.withDefaultNamespace("entity/banner_base"));
	public static final Material SHIELD_BASE = new Material(Sheets.SHIELD_SHEET, ResourceLocation.withDefaultNamespace("entity/shield_base"));
	public static final Material NO_PATTERN_SHIELD = new Material(Sheets.SHIELD_SHEET, ResourceLocation.withDefaultNamespace("entity/shield_base_nopattern"));
	public static final int DESTROY_STAGE_COUNT = 10;
	public static final List<ResourceLocation> DESTROY_STAGES = (List<ResourceLocation>)IntStream.range(0, 10)
		.mapToObj(i -> ResourceLocation.withDefaultNamespace("block/destroy_stage_" + i))
		.collect(Collectors.toList());
	public static final List<ResourceLocation> BREAKING_LOCATIONS = (List<ResourceLocation>)DESTROY_STAGES.stream()
		.map(resourceLocation -> resourceLocation.withPath((UnaryOperator<String>)(string -> "textures/" + string + ".png")))
		.collect(Collectors.toList());
	public static final List<RenderType> DESTROY_TYPES = (List<RenderType>)BREAKING_LOCATIONS.stream().map(RenderType::crumbling).collect(Collectors.toList());
	static final Logger LOGGER = LogUtils.getLogger();
	private final EntityModelSet entityModelSet;
	private final Map<BlockState, BlockStateModel.UnbakedRoot> unbakedBlockStateModels;
	private final Map<ResourceLocation, ClientItem> clientInfos;
	final Map<ResourceLocation, ResolvedModel> resolvedModels;
	final ResolvedModel missingModel;

	public ModelBakery(
		EntityModelSet entityModelSet,
		Map<BlockState, BlockStateModel.UnbakedRoot> unbakedBlockStateModels,
		Map<ResourceLocation, ClientItem> clientInfos,
		Map<ResourceLocation, ResolvedModel> resolvedModels,
		ResolvedModel missingModel
	) {
		this.entityModelSet = entityModelSet;
		this.unbakedBlockStateModels = unbakedBlockStateModels;
		this.clientInfos = clientInfos;
		this.resolvedModels = resolvedModels;
		this.missingModel = missingModel;
	}

	public CompletableFuture<ModelBakery.BakingResult> bakeModels(SpriteGetter sprites, Executor executor) {
		ModelBakery.MissingModels missingModels = ModelBakery.MissingModels.bake(this.missingModel, sprites);
		ModelBakery.ModelBakerImpl modelBakerImpl = new ModelBakery.ModelBakerImpl(sprites);
		CompletableFuture<Map<BlockState, BlockStateModel>> completableFuture = ParallelMapTransform.schedule(
			this.unbakedBlockStateModels, (blockState, unbakedRoot) -> {
				try {
					return unbakedRoot.bake(blockState, modelBakerImpl);
				} catch (Exception var4x) {
					LOGGER.warn("Unable to bake model: '{}': {}", blockState, var4x);
					return null;
				}
			}, executor
		);
		CompletableFuture<Map<ResourceLocation, ItemModel>> completableFuture2 = ParallelMapTransform.schedule(this.clientInfos, (resourceLocation, clientItem) -> {
			try {
				return clientItem.model().bake(new ItemModel.BakingContext(modelBakerImpl, this.entityModelSet, missingModels.item, clientItem.registrySwapper()));
			} catch (Exception var6x) {
				LOGGER.warn("Unable to bake item model: '{}'", resourceLocation, var6x);
				return null;
			}
		}, executor);
		Map<ResourceLocation, ClientItem.Properties> map = new HashMap(this.clientInfos.size());
		this.clientInfos.forEach((resourceLocation, clientItem) -> {
			ClientItem.Properties properties = clientItem.properties();
			if (!properties.equals(ClientItem.Properties.DEFAULT)) {
				map.put(resourceLocation, properties);
			}
		});
		return completableFuture.thenCombine(completableFuture2, (map2, map3) -> new ModelBakery.BakingResult(missingModels, map2, map3, map));
	}

	@Environment(EnvType.CLIENT)
	public record BakingResult(
		ModelBakery.MissingModels missingModels,
		Map<BlockState, BlockStateModel> blockStateModels,
		Map<ResourceLocation, ItemModel> itemStackModels,
		Map<ResourceLocation, ClientItem.Properties> itemProperties
	) {
	}

	@Environment(EnvType.CLIENT)
	public record MissingModels(BlockStateModel block, ItemModel item) {

		public static ModelBakery.MissingModels bake(ResolvedModel model, SpriteGetter sprites) {
			ModelBaker modelBaker = new 1(sprites);
			TextureSlots textureSlots = model.getTopTextureSlots();
			boolean bl = model.getTopAmbientOcclusion();
			boolean bl2 = model.getTopGuiLight().lightLikeBlock();
			ItemTransforms itemTransforms = model.getTopTransforms();
			QuadCollection quadCollection = model.bakeTopGeometry(textureSlots, modelBaker, BlockModelRotation.X0_Y0);
			TextureAtlasSprite textureAtlasSprite = model.resolveParticleSprite(textureSlots, modelBaker);
			BlockStateModel blockStateModel = new SingleVariant(new SimpleModelWrapper(quadCollection, bl, textureAtlasSprite));
			ItemModel itemModel = new MissingItemModel(quadCollection.getAll(), new ModelRenderProperties(bl2, textureAtlasSprite, itemTransforms));
			return new ModelBakery.MissingModels(blockStateModel, itemModel);
		}
	}

	@Environment(EnvType.CLIENT)
	class ModelBakerImpl implements ModelBaker {
		private final SpriteGetter sprites;
		private final Map<ModelBaker.SharedOperationKey<Object>, Object> operationCache = new ConcurrentHashMap();
		private final Function<ModelBaker.SharedOperationKey<Object>, Object> cacheComputeFunction = sharedOperationKey -> sharedOperationKey.compute(this);

		ModelBakerImpl(final SpriteGetter sprites) {
			this.sprites = sprites;
		}

		@Override
		public SpriteGetter sprites() {
			return this.sprites;
		}

		@Override
		public ResolvedModel getModel(ResourceLocation modelLocation) {
			ResolvedModel resolvedModel = (ResolvedModel)ModelBakery.this.resolvedModels.get(modelLocation);
			if (resolvedModel == null) {
				ModelBakery.LOGGER.warn("Requested a model that was not discovered previously: {}", modelLocation);
				return ModelBakery.this.missingModel;
			} else {
				return resolvedModel;
			}
		}

		@Override
		public <T> T compute(ModelBaker.SharedOperationKey<T> key) {
			return (T)this.operationCache.computeIfAbsent(key, this.cacheComputeFunction);
		}
	}
}