package net.minecraft.client.resources.model;

import com.google.common.collect.Maps;
import com.google.gson.JsonObject;
import com.mojang.logging.LogUtils;
import java.io.Reader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.Function;
import java.util.stream.Stream;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.Util;
import net.minecraft.client.renderer.block.BlockModelShaper;
import net.minecraft.client.renderer.block.model.BlockModelDefinition;
import net.minecraft.client.renderer.block.model.UnbakedBlockStateModel;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.resources.FileToIdConverter;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.packs.resources.Resource;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.util.GsonHelper;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.Blocks;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.block.state.StateDefinition;
import net.minecraft.world.level.block.state.StateDefinition.Builder;
import net.minecraft.world.level.block.state.properties.BooleanProperty;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class BlockStateModelLoader {
	private static final Logger LOGGER = LogUtils.getLogger();
	private static final FileToIdConverter BLOCKSTATE_LISTER = FileToIdConverter.json("blockstates");
	private static final String FRAME_MAP_PROPERTY = "map";
	private static final String FRAME_MAP_PROPERTY_TRUE = "map=true";
	private static final String FRAME_MAP_PROPERTY_FALSE = "map=false";
	private static final StateDefinition<Block, BlockState> ITEM_FRAME_FAKE_DEFINITION = new Builder<Block, BlockState>(Blocks.AIR)
		.add(BooleanProperty.create("map"))
		.create(Block::defaultBlockState, BlockState::new);
	private static final ResourceLocation GLOW_ITEM_FRAME_LOCATION = ResourceLocation.withDefaultNamespace("glow_item_frame");
	private static final ResourceLocation ITEM_FRAME_LOCATION = ResourceLocation.withDefaultNamespace("item_frame");
	private static final Map<ResourceLocation, StateDefinition<Block, BlockState>> STATIC_DEFINITIONS = Map.of(
		ITEM_FRAME_LOCATION, ITEM_FRAME_FAKE_DEFINITION, GLOW_ITEM_FRAME_LOCATION, ITEM_FRAME_FAKE_DEFINITION
	);
	public static final ModelResourceLocation GLOW_MAP_FRAME_LOCATION = new ModelResourceLocation(GLOW_ITEM_FRAME_LOCATION, "map=true");
	public static final ModelResourceLocation GLOW_FRAME_LOCATION = new ModelResourceLocation(GLOW_ITEM_FRAME_LOCATION, "map=false");
	public static final ModelResourceLocation MAP_FRAME_LOCATION = new ModelResourceLocation(ITEM_FRAME_LOCATION, "map=true");
	public static final ModelResourceLocation FRAME_LOCATION = new ModelResourceLocation(ITEM_FRAME_LOCATION, "map=false");

	private static Function<ResourceLocation, StateDefinition<Block, BlockState>> definitionLocationToBlockMapper() {
		Map<ResourceLocation, StateDefinition<Block, BlockState>> map = new HashMap(STATIC_DEFINITIONS);

		for (Block block : BuiltInRegistries.BLOCK) {
			map.put(block.builtInRegistryHolder().key().location(), block.getStateDefinition());
		}

		return map::get;
	}

	public static CompletableFuture<BlockStateModelLoader.LoadedModels> loadBlockStates(UnbakedModel model, ResourceManager resourceManager, Executor executor) {
		Function<ResourceLocation, StateDefinition<Block, BlockState>> function = definitionLocationToBlockMapper();
		return CompletableFuture.supplyAsync(() -> BLOCKSTATE_LISTER.listMatchingResourceStacks(resourceManager), executor).thenCompose(map -> {
			List<CompletableFuture<BlockStateModelLoader.LoadedModels>> list = new ArrayList(map.size());

			for (Entry<ResourceLocation, List<Resource>> entry : map.entrySet()) {
				list.add(CompletableFuture.supplyAsync(() -> {
					ResourceLocation resourceLocation = BLOCKSTATE_LISTER.fileToId((ResourceLocation)entry.getKey());
					StateDefinition<Block, BlockState> stateDefinition = (StateDefinition<Block, BlockState>)function.apply(resourceLocation);
					if (stateDefinition == null) {
						LOGGER.debug("Discovered unknown block state definition {}, ignoring", resourceLocation);
						return null;
					} else {
						List<Resource> listx = (List<Resource>)entry.getValue();
						List<BlockStateModelLoader.LoadedBlockModelDefinition> list2 = new ArrayList(listx.size());

						for (Resource resource : listx) {
							try {
								Reader reader = resource.openAsReader();

								try {
									JsonObject jsonObject = GsonHelper.parse(reader);
									BlockModelDefinition blockModelDefinition = BlockModelDefinition.fromJsonElement(jsonObject);
									list2.add(new BlockStateModelLoader.LoadedBlockModelDefinition(resource.sourcePackId(), blockModelDefinition));
								} catch (Throwable var14) {
									if (reader != null) {
										try {
											reader.close();
										} catch (Throwable var13) {
											var14.addSuppressed(var13);
										}
									}

									throw var14;
								}

								if (reader != null) {
									reader.close();
								}
							} catch (Exception var15) {
								LOGGER.error("Failed to load blockstate definition {} from pack {}", resourceLocation, resource.sourcePackId(), var15);
							}
						}

						try {
							return loadBlockStateDefinitionStack(resourceLocation, stateDefinition, list2, model);
						} catch (Exception var12) {
							LOGGER.error("Failed to load blockstate definition {}", resourceLocation, var12);
							return null;
						}
					}
				}, executor));
			}

			return Util.sequence(list).thenApply(listx -> {
				Map<ModelResourceLocation, BlockStateModelLoader.LoadedModel> mapx = new HashMap();

				for (BlockStateModelLoader.LoadedModels loadedModels : listx) {
					if (loadedModels != null) {
						mapx.putAll(loadedModels.models());
					}
				}

				return new BlockStateModelLoader.LoadedModels(mapx);
			});
		});
	}

	private static BlockStateModelLoader.LoadedModels loadBlockStateDefinitionStack(
		ResourceLocation id,
		StateDefinition<Block, BlockState> stateDefinition,
		List<BlockStateModelLoader.LoadedBlockModelDefinition> modelDefinitions,
		UnbakedModel model
	) {
		Map<ModelResourceLocation, BlockStateModelLoader.LoadedModel> map = new HashMap();

		for (BlockStateModelLoader.LoadedBlockModelDefinition loadedBlockModelDefinition : modelDefinitions) {
			loadedBlockModelDefinition.contents
				.instantiate(stateDefinition, id + "/" + loadedBlockModelDefinition.source)
				.forEach((blockState, unbakedBlockStateModel) -> {
					ModelResourceLocation modelResourceLocation = BlockModelShaper.stateToModelLocation(id, blockState);
					map.put(modelResourceLocation, new BlockStateModelLoader.LoadedModel(blockState, unbakedBlockStateModel));
				});
		}

		return new BlockStateModelLoader.LoadedModels(map);
	}

	@Environment(EnvType.CLIENT)
	record LoadedBlockModelDefinition(String source, BlockModelDefinition contents) {
	}

	@Environment(EnvType.CLIENT)
	public record LoadedModel(BlockState state, UnbakedBlockStateModel model) {
	}

	@Environment(EnvType.CLIENT)
	public record LoadedModels(Map<ModelResourceLocation, BlockStateModelLoader.LoadedModel> models) {
		public Stream<ResolvableModel> forResolving() {
			return this.models.values().stream().map(BlockStateModelLoader.LoadedModel::model);
		}

		public Map<ModelResourceLocation, UnbakedBlockStateModel> plainModels() {
			return Maps.transformValues(this.models, BlockStateModelLoader.LoadedModel::model);
		}
	}
}