package net.minecraft.client.resources.model;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.mojang.datafixers.util.Pair;
import com.mojang.logging.LogUtils;
import it.unimi.dsi.fastutil.objects.Object2IntMap;
import it.unimi.dsi.fastutil.objects.Object2IntMaps;
import java.io.Reader;
import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.Map.Entry;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.Util;
import net.minecraft.client.color.block.BlockColors;
import net.minecraft.client.model.geom.EntityModelSet;
import net.minecraft.client.renderer.Sheets;
import net.minecraft.client.renderer.SpecialBlockModelRenderer;
import net.minecraft.client.renderer.block.BlockModelShaper;
import net.minecraft.client.renderer.block.model.BlockModel;
import net.minecraft.client.renderer.item.ClientItem;
import net.minecraft.client.renderer.item.ItemModel;
import net.minecraft.client.renderer.texture.TextureAtlas;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.renderer.texture.TextureManager;
import net.minecraft.client.resources.model.AtlasSet.StitchResult;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.resources.FileToIdConverter;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.packs.resources.PreparableReloadListener;
import net.minecraft.server.packs.resources.Resource;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.server.packs.resources.PreparableReloadListener.PreparationBarrier;
import net.minecraft.util.profiling.Profiler;
import net.minecraft.util.profiling.ProfilerFiller;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.material.FluidState;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class ModelManager implements PreparableReloadListener, AutoCloseable {
	private static final Logger LOGGER = LogUtils.getLogger();
	private static final FileToIdConverter MODEL_LISTER = FileToIdConverter.json("models");
	private static final Map<ResourceLocation, ResourceLocation> VANILLA_ATLASES = Map.of(
		Sheets.BANNER_SHEET,
		ResourceLocation.withDefaultNamespace("banner_patterns"),
		Sheets.BED_SHEET,
		ResourceLocation.withDefaultNamespace("beds"),
		Sheets.CHEST_SHEET,
		ResourceLocation.withDefaultNamespace("chests"),
		Sheets.SHIELD_SHEET,
		ResourceLocation.withDefaultNamespace("shield_patterns"),
		Sheets.SIGN_SHEET,
		ResourceLocation.withDefaultNamespace("signs"),
		Sheets.SHULKER_SHEET,
		ResourceLocation.withDefaultNamespace("shulker_boxes"),
		Sheets.ARMOR_TRIMS_SHEET,
		ResourceLocation.withDefaultNamespace("armor_trims"),
		Sheets.DECORATED_POT_SHEET,
		ResourceLocation.withDefaultNamespace("decorated_pot"),
		TextureAtlas.LOCATION_BLOCKS,
		ResourceLocation.withDefaultNamespace("blocks")
	);
	private Map<ModelResourceLocation, BakedModel> bakedBlockStateModels = Map.of();
	private Map<ResourceLocation, ItemModel> bakedItemStackModels = Map.of();
	private Map<ResourceLocation, ClientItem.Properties> itemProperties = Map.of();
	private final AtlasSet atlases;
	private final BlockModelShaper blockModelShaper;
	private final BlockColors blockColors;
	private EntityModelSet entityModelSet = EntityModelSet.EMPTY;
	private SpecialBlockModelRenderer specialBlockModelRenderer = SpecialBlockModelRenderer.EMPTY;
	private int maxMipmapLevels;
	private BakedModel missingModel;
	private ItemModel missingItemModel;
	private Object2IntMap<BlockState> modelGroups = Object2IntMaps.emptyMap();

	public ModelManager(TextureManager textureManager, BlockColors blockColors, int maxMipmapLevels) {
		this.blockColors = blockColors;
		this.maxMipmapLevels = maxMipmapLevels;
		this.blockModelShaper = new BlockModelShaper(this);
		this.atlases = new AtlasSet(VANILLA_ATLASES, textureManager);
	}

	public BakedModel getModel(ModelResourceLocation modelLocation) {
		return (BakedModel)this.bakedBlockStateModels.getOrDefault(modelLocation, this.missingModel);
	}

	public BakedModel getMissingModel() {
		return this.missingModel;
	}

	public ItemModel getItemModel(ResourceLocation modelLocation) {
		return (ItemModel)this.bakedItemStackModels.getOrDefault(modelLocation, this.missingItemModel);
	}

	public ClientItem.Properties getItemProperties(ResourceLocation itemId) {
		return (ClientItem.Properties)this.itemProperties.getOrDefault(itemId, ClientItem.Properties.DEFAULT);
	}

	public BlockModelShaper getBlockModelShaper() {
		return this.blockModelShaper;
	}

	@Override
	public final CompletableFuture<Void> reload(PreparationBarrier barrier, ResourceManager manager, Executor backgroundExecutor, Executor gameExecutor) {
		UnbakedModel unbakedModel = MissingBlockModel.missingModel();
		CompletableFuture<EntityModelSet> completableFuture = CompletableFuture.supplyAsync(EntityModelSet::vanilla, backgroundExecutor);
		CompletableFuture<SpecialBlockModelRenderer> completableFuture2 = completableFuture.thenApplyAsync(SpecialBlockModelRenderer::vanilla, backgroundExecutor);
		CompletableFuture<Map<ResourceLocation, UnbakedModel>> completableFuture3 = loadBlockModels(manager, backgroundExecutor);
		CompletableFuture<BlockStateModelLoader.LoadedModels> completableFuture4 = BlockStateModelLoader.loadBlockStates(unbakedModel, manager, backgroundExecutor);
		CompletableFuture<ClientItemInfoLoader.LoadedClientInfos> completableFuture5 = ClientItemInfoLoader.scheduleLoad(manager, backgroundExecutor);
		CompletableFuture<ModelDiscovery> completableFuture6 = CompletableFuture.allOf(completableFuture3, completableFuture4, completableFuture5)
			.thenApplyAsync(
				void_ -> discoverModelDependencies(
					unbakedModel,
					(Map<ResourceLocation, UnbakedModel>)completableFuture3.join(),
					(BlockStateModelLoader.LoadedModels)completableFuture4.join(),
					(ClientItemInfoLoader.LoadedClientInfos)completableFuture5.join()
				),
				backgroundExecutor
			);
		CompletableFuture<Object2IntMap<BlockState>> completableFuture7 = completableFuture4.thenApplyAsync(
			loadedModels -> buildModelGroups(this.blockColors, loadedModels), backgroundExecutor
		);
		Map<ResourceLocation, CompletableFuture<StitchResult>> map = this.atlases.scheduleLoad(manager, this.maxMipmapLevels, backgroundExecutor);
		return CompletableFuture.allOf(
				(CompletableFuture[])Stream.concat(
						map.values().stream(), Stream.of(completableFuture6, completableFuture7, completableFuture4, completableFuture5, completableFuture, completableFuture2)
					)
					.toArray(CompletableFuture[]::new)
			)
			.thenApplyAsync(
				void_ -> {
					Map<ResourceLocation, StitchResult> map2 = (Map<ResourceLocation, StitchResult>)map.entrySet()
						.stream()
						.collect(Collectors.toMap(Entry::getKey, entry -> (StitchResult)((CompletableFuture)entry.getValue()).join()));
					ModelDiscovery modelDiscovery = (ModelDiscovery)completableFuture6.join();
					Object2IntMap<BlockState> object2IntMap = (Object2IntMap<BlockState>)completableFuture7.join();
					Set<ResourceLocation> set = modelDiscovery.getUnreferencedModels();
					if (!set.isEmpty()) {
						LOGGER.debug("Unreferenced models: \n{}", set.stream().sorted().map(resourceLocation -> "\t" + resourceLocation + "\n").collect(Collectors.joining()));
					}

					ModelBakery modelBakery = new ModelBakery(
						(EntityModelSet)completableFuture.join(),
						((BlockStateModelLoader.LoadedModels)completableFuture4.join()).plainModels(),
						((ClientItemInfoLoader.LoadedClientInfos)completableFuture5.join()).contents(),
						modelDiscovery.getReferencedModels(),
						unbakedModel
					);
					return loadModels(
						Profiler.get(), map2, modelBakery, object2IntMap, (EntityModelSet)completableFuture.join(), (SpecialBlockModelRenderer)completableFuture2.join()
					);
				},
				backgroundExecutor
			)
			.thenCompose(reloadState -> reloadState.readyForUpload.thenApply(void_ -> reloadState))
			.thenCompose(barrier::wait)
			.thenAcceptAsync(reloadState -> this.apply(reloadState, Profiler.get()), gameExecutor);
	}

	private static CompletableFuture<Map<ResourceLocation, UnbakedModel>> loadBlockModels(ResourceManager resourceManager, Executor executor) {
		return CompletableFuture.supplyAsync(() -> MODEL_LISTER.listMatchingResources(resourceManager), executor)
			.thenCompose(
				map -> {
					List<CompletableFuture<Pair<ResourceLocation, BlockModel>>> list = new ArrayList(map.size());

					for (Entry<ResourceLocation, Resource> entry : map.entrySet()) {
						list.add(CompletableFuture.supplyAsync(() -> {
							ResourceLocation resourceLocation = MODEL_LISTER.fileToId((ResourceLocation)entry.getKey());

							try {
								Reader reader = ((Resource)entry.getValue()).openAsReader();

								Pair var3;
								try {
									var3 = Pair.of(resourceLocation, BlockModel.fromStream(reader));
								} catch (Throwable var6) {
									if (reader != null) {
										try {
											reader.close();
										} catch (Throwable var5) {
											var6.addSuppressed(var5);
										}
									}

									throw var6;
								}

								if (reader != null) {
									reader.close();
								}

								return var3;
							} catch (Exception var7) {
								LOGGER.error("Failed to load model {}", entry.getKey(), var7);
								return null;
							}
						}, executor));
					}

					return Util.sequence(list)
						.thenApply(listx -> (Map)listx.stream().filter(Objects::nonNull).collect(Collectors.toUnmodifiableMap(Pair::getFirst, Pair::getSecond)));
				}
			);
	}

	private static ModelDiscovery discoverModelDependencies(
		UnbakedModel missingModel,
		Map<ResourceLocation, UnbakedModel> inputModels,
		BlockStateModelLoader.LoadedModels loadedModels,
		ClientItemInfoLoader.LoadedClientInfos loadedClientInfos
	) {
		ModelDiscovery modelDiscovery = new ModelDiscovery(inputModels, missingModel);
		loadedModels.forResolving().forEach(modelDiscovery::addRoot);
		loadedClientInfos.contents().values().forEach(clientItem -> modelDiscovery.addRoot(clientItem.model()));
		modelDiscovery.registerSpecialModels();
		modelDiscovery.discoverDependencies();
		return modelDiscovery;
	}

	private static ModelManager.ReloadState loadModels(
		ProfilerFiller profiler,
		Map<ResourceLocation, StitchResult> atlasPreperations,
		ModelBakery modelBakery,
		Object2IntMap<BlockState> modelGroups,
		EntityModelSet entityModelSet,
		SpecialBlockModelRenderer specialBlockModelRenderer
	) {
		profiler.push("baking");
		final Multimap<String, Material> multimap = HashMultimap.create();
		final Multimap<String, String> multimap2 = HashMultimap.create();
		final TextureAtlasSprite textureAtlasSprite = ((StitchResult)atlasPreperations.get(TextureAtlas.LOCATION_BLOCKS)).missing();
		ModelBakery.BakingResult bakingResult = modelBakery.bakeModels(new ModelBakery.TextureGetter() {
			@Override
			public TextureAtlasSprite get(ModelDebugName name, Material material) {
				StitchResult stitchResult = (StitchResult)atlasPreperations.get(material.atlasLocation());
				TextureAtlasSprite textureAtlasSpritex = stitchResult.getSprite(material.texture());
				if (textureAtlasSpritex != null) {
					return textureAtlasSpritex;
				} else {
					multimap.put((String)name.get(), material);
					return stitchResult.missing();
				}
			}

			@Override
			public TextureAtlasSprite reportMissingReference(ModelDebugName name, String reference) {
				multimap2.put((String)name.get(), reference);
				return textureAtlasSprite;
			}
		});
		multimap.asMap()
			.forEach(
				(string, collection) -> LOGGER.warn(
					"Missing textures in model {}:\n{}",
					string,
					collection.stream()
						.sorted(Material.COMPARATOR)
						.map(material -> "    " + material.atlasLocation() + ":" + material.texture())
						.collect(Collectors.joining("\n"))
				)
			);
		multimap2.asMap()
			.forEach(
				(string, collection) -> LOGGER.warn(
					"Missing texture references in model {}:\n{}", string, collection.stream().sorted().map(stringx -> "    " + stringx).collect(Collectors.joining("\n"))
				)
			);
		profiler.popPush("dispatch");
		Map<BlockState, BakedModel> map = createBlockStateToModelDispatch(bakingResult.blockStateModels(), bakingResult.missingModel());
		CompletableFuture<Void> completableFuture = CompletableFuture.allOf(
			(CompletableFuture[])atlasPreperations.values().stream().map(StitchResult::readyForUpload).toArray(CompletableFuture[]::new)
		);
		profiler.pop();
		return new ModelManager.ReloadState(bakingResult, modelGroups, map, atlasPreperations, entityModelSet, specialBlockModelRenderer, completableFuture);
	}

	private static Map<BlockState, BakedModel> createBlockStateToModelDispatch(Map<ModelResourceLocation, BakedModel> blockStateModels, BakedModel missingModel) {
		Map<BlockState, BakedModel> map = new IdentityHashMap();

		for (Block block : BuiltInRegistries.BLOCK) {
			block.getStateDefinition().getPossibleStates().forEach(blockState -> {
				ResourceLocation resourceLocation = blockState.getBlock().builtInRegistryHolder().key().location();
				ModelResourceLocation modelResourceLocation = BlockModelShaper.stateToModelLocation(resourceLocation, blockState);
				BakedModel bakedModel2 = (BakedModel)blockStateModels.get(modelResourceLocation);
				if (bakedModel2 == null) {
					LOGGER.warn("Missing model for variant: '{}'", modelResourceLocation);
					map.putIfAbsent(blockState, missingModel);
				} else {
					map.put(blockState, bakedModel2);
				}
			});
		}

		return map;
	}

	private static Object2IntMap<BlockState> buildModelGroups(BlockColors blockColors, BlockStateModelLoader.LoadedModels loadedModels) {
		return ModelGroupCollector.build(blockColors, loadedModels);
	}

	private void apply(ModelManager.ReloadState reloadState, ProfilerFiller profiler) {
		profiler.push("upload");
		reloadState.atlasPreparations.values().forEach(StitchResult::upload);
		ModelBakery.BakingResult bakingResult = reloadState.bakedModels;
		this.bakedBlockStateModels = bakingResult.blockStateModels();
		this.bakedItemStackModels = bakingResult.itemStackModels();
		this.itemProperties = bakingResult.itemProperties();
		this.modelGroups = reloadState.modelGroups;
		this.missingModel = bakingResult.missingModel();
		this.missingItemModel = bakingResult.missingItemModel();
		profiler.popPush("cache");
		this.blockModelShaper.replaceCache(reloadState.modelCache);
		this.specialBlockModelRenderer = reloadState.specialBlockModelRenderer;
		this.entityModelSet = reloadState.entityModelSet;
		profiler.pop();
	}

	public boolean requiresRender(BlockState oldState, BlockState newState) {
		if (oldState == newState) {
			return false;
		} else {
			int i = this.modelGroups.getInt(oldState);
			if (i != -1) {
				int j = this.modelGroups.getInt(newState);
				if (i == j) {
					FluidState fluidState = oldState.getFluidState();
					FluidState fluidState2 = newState.getFluidState();
					return fluidState != fluidState2;
				}
			}

			return true;
		}
	}

	public TextureAtlas getAtlas(ResourceLocation location) {
		return this.atlases.getAtlas(location);
	}

	public void close() {
		this.atlases.close();
	}

	public void updateMaxMipLevel(int level) {
		this.maxMipmapLevels = level;
	}

	public Supplier<SpecialBlockModelRenderer> specialBlockModelRenderer() {
		return () -> this.specialBlockModelRenderer;
	}

	public Supplier<EntityModelSet> entityModels() {
		return () -> this.entityModelSet;
	}

	@Environment(EnvType.CLIENT)
	record ReloadState(
		ModelBakery.BakingResult bakedModels,
		Object2IntMap<BlockState> modelGroups,
		Map<BlockState, BakedModel> modelCache,
		Map<ResourceLocation, StitchResult> atlasPreparations,
		EntityModelSet entityModelSet,
		SpecialBlockModelRenderer specialBlockModelRenderer,
		CompletableFuture<Void> readyForUpload
	) {
	}
}