/*
 * Decompiled with CFR 0.152.
 * 
 * Could not load the following classes:
 *  com.google.common.collect.HashMultimap
 *  com.google.common.collect.Multimap
 *  com.google.common.collect.Multimaps
 *  com.google.common.collect.Sets
 *  com.google.common.collect.Sets$SetView
 *  com.mojang.datafixers.util.Pair
 *  com.mojang.logging.LogUtils
 *  it.unimi.dsi.fastutil.objects.Object2IntMap
 *  it.unimi.dsi.fastutil.objects.Object2IntMaps
 *  org.jspecify.annotations.Nullable
 *  org.slf4j.Logger
 */
package net.minecraft.client.resources.model;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Sets;
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.BufferedReader;
import java.io.Reader;
import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Executor;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import net.minecraft.client.color.block.BlockColors;
import net.minecraft.client.model.geom.EntityModelSet;
import net.minecraft.client.renderer.PlayerSkinRenderCache;
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.block.model.BlockStateModel;
import net.minecraft.client.renderer.block.model.ItemModelGenerator;
import net.minecraft.client.renderer.item.ClientItem;
import net.minecraft.client.renderer.item.ItemModel;
import net.minecraft.client.renderer.special.SpecialModelRenderer;
import net.minecraft.client.renderer.texture.SpriteLoader;
import net.minecraft.client.renderer.texture.TextureAtlas;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.resources.model.AtlasManager;
import net.minecraft.client.resources.model.BlockStateModelLoader;
import net.minecraft.client.resources.model.ClientItemInfoLoader;
import net.minecraft.client.resources.model.Material;
import net.minecraft.client.resources.model.MissingBlockModel;
import net.minecraft.client.resources.model.ModelBakery;
import net.minecraft.client.resources.model.ModelDebugName;
import net.minecraft.client.resources.model.ModelDiscovery;
import net.minecraft.client.resources.model.ModelGroupCollector;
import net.minecraft.client.resources.model.ResolvedModel;
import net.minecraft.client.resources.model.SpriteGetter;
import net.minecraft.client.resources.model.UnbakedModel;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.data.AtlasIds;
import net.minecraft.resources.FileToIdConverter;
import net.minecraft.resources.Identifier;
import net.minecraft.server.packs.resources.PreparableReloadListener;
import net.minecraft.server.packs.resources.Resource;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.util.Util;
import net.minecraft.util.profiling.Profiler;
import net.minecraft.util.profiling.Zone;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.material.FluidState;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;

public class ModelManager
implements PreparableReloadListener {
    public static final Identifier BLOCK_OR_ITEM = Identifier.withDefaultNamespace("block_or_item");
    private static final Logger LOGGER = LogUtils.getLogger();
    private static final FileToIdConverter MODEL_LISTER = FileToIdConverter.json("models");
    private Map<Identifier, ItemModel> bakedItemStackModels = Map.of();
    private Map<Identifier, ClientItem.Properties> itemProperties = Map.of();
    private final AtlasManager atlasManager;
    private final PlayerSkinRenderCache playerSkinRenderCache;
    private final BlockModelShaper blockModelShaper;
    private final BlockColors blockColors;
    private EntityModelSet entityModelSet = EntityModelSet.EMPTY;
    private SpecialBlockModelRenderer specialBlockModelRenderer = SpecialBlockModelRenderer.EMPTY;
    private ModelBakery.MissingModels missingModels;
    private Object2IntMap<BlockState> modelGroups = Object2IntMaps.emptyMap();

    public ModelManager(BlockColors blockColors, AtlasManager atlasManager, PlayerSkinRenderCache playerSkinRenderCache) {
        this.blockColors = blockColors;
        this.atlasManager = atlasManager;
        this.playerSkinRenderCache = playerSkinRenderCache;
        this.blockModelShaper = new BlockModelShaper(this);
    }

    public BlockStateModel getMissingBlockStateModel() {
        return this.missingModels.block();
    }

    public ItemModel getItemModel(Identifier id) {
        return this.bakedItemStackModels.getOrDefault(id, this.missingModels.item());
    }

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

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

    @Override
    public final CompletableFuture<Void> reload(PreparableReloadListener.SharedState currentReload, Executor taskExecutor, PreparableReloadListener.PreparationBarrier preparationBarrier, Executor reloadExecutor) {
        ResourceManager manager = currentReload.resourceManager();
        CompletableFuture<EntityModelSet> entityModelSet = CompletableFuture.supplyAsync(EntityModelSet::vanilla, taskExecutor);
        CompletionStage specialBlockModelRenderer = entityModelSet.thenApplyAsync(entityModels -> SpecialBlockModelRenderer.vanilla(new SpecialModelRenderer.BakingContext.Simple((EntityModelSet)entityModels, this.atlasManager, this.playerSkinRenderCache)), taskExecutor);
        CompletableFuture<Map<Identifier, UnbakedModel>> modelCache = ModelManager.loadBlockModels(manager, taskExecutor);
        CompletableFuture<BlockStateModelLoader.LoadedModels> blockStateModels = BlockStateModelLoader.loadBlockStates(manager, taskExecutor);
        CompletableFuture<ClientItemInfoLoader.LoadedClientInfos> itemStackModels = ClientItemInfoLoader.scheduleLoad(manager, taskExecutor);
        CompletionStage modelDiscovery = CompletableFuture.allOf(modelCache, blockStateModels, itemStackModels).thenApplyAsync(unused -> ModelManager.discoverModelDependencies((Map)modelCache.join(), (BlockStateModelLoader.LoadedModels)blockStateModels.join(), (ClientItemInfoLoader.LoadedClientInfos)itemStackModels.join()), taskExecutor);
        CompletionStage modelGroups = blockStateModels.thenApplyAsync(models -> ModelManager.buildModelGroups(this.blockColors, models), taskExecutor);
        AtlasManager.PendingStitchResults pendingStitches = currentReload.get(AtlasManager.PENDING_STITCH);
        CompletableFuture<SpriteLoader.Preparations> pendingBlockAtlasSprites = pendingStitches.get(AtlasIds.BLOCKS);
        CompletableFuture<SpriteLoader.Preparations> pendingItemAtlasSprites = pendingStitches.get(AtlasIds.ITEMS);
        return ((CompletableFuture)((CompletableFuture)CompletableFuture.allOf(new CompletableFuture[]{pendingBlockAtlasSprites, pendingItemAtlasSprites, modelDiscovery, modelGroups, blockStateModels, itemStackModels, entityModelSet, specialBlockModelRenderer, modelCache}).thenComposeAsync(arg_0 -> this.lambda$reload$4(pendingBlockAtlasSprites, pendingItemAtlasSprites, (CompletableFuture)modelDiscovery, (CompletableFuture)modelGroups, modelCache, entityModelSet, blockStateModels, itemStackModels, (CompletableFuture)specialBlockModelRenderer, taskExecutor, arg_0), taskExecutor)).thenCompose(preparationBarrier::wait)).thenAcceptAsync(this::apply, reloadExecutor);
    }

    private static CompletableFuture<Map<Identifier, UnbakedModel>> loadBlockModels(ResourceManager manager, Executor executor) {
        return CompletableFuture.supplyAsync(() -> MODEL_LISTER.listMatchingResources(manager), executor).thenCompose(resources -> {
            ArrayList<CompletableFuture<@Nullable Pair>> result = new ArrayList<CompletableFuture<Pair>>(resources.size());
            for (Map.Entry resource : resources.entrySet()) {
                result.add(CompletableFuture.supplyAsync(() -> {
                    Pair pair;
                    block8: {
                        Identifier modelId = MODEL_LISTER.fileToId((Identifier)resource.getKey());
                        @Nullable BufferedReader reader = ((Resource)resource.getValue()).openAsReader();
                        try {
                            pair = Pair.of((Object)modelId, (Object)BlockModel.fromStream(reader));
                            if (reader == null) break block8;
                        }
                        catch (Throwable throwable) {
                            try {
                                if (reader != null) {
                                    try {
                                        ((Reader)reader).close();
                                    }
                                    catch (Throwable throwable2) {
                                        throwable.addSuppressed(throwable2);
                                    }
                                }
                                throw throwable;
                            }
                            catch (Exception e) {
                                LOGGER.error("Failed to load model {}", resource.getKey(), (Object)e);
                                return null;
                            }
                        }
                        ((Reader)reader).close();
                    }
                    return pair;
                }, executor));
            }
            return Util.sequence(result).thenApply(pairs -> pairs.stream().filter(Objects::nonNull).collect(Collectors.toUnmodifiableMap(Pair::getFirst, Pair::getSecond)));
        });
    }

    private static ResolvedModels discoverModelDependencies(Map<Identifier, UnbakedModel> allModels, BlockStateModelLoader.LoadedModels blockStateModels, ClientItemInfoLoader.LoadedClientInfos itemInfos) {
        try (Zone ignored = Profiler.get().zone("dependencies");){
            ModelDiscovery result = new ModelDiscovery(allModels, MissingBlockModel.missingModel());
            result.addSpecialModel(ItemModelGenerator.GENERATED_ITEM_MODEL_ID, new ItemModelGenerator());
            blockStateModels.models().values().forEach(result::addRoot);
            itemInfos.contents().values().forEach(info -> result.addRoot(info.model()));
            ResolvedModels resolvedModels = new ResolvedModels(result.missingModel(), result.resolve());
            return resolvedModels;
        }
    }

    private static CompletableFuture<ReloadState> loadModels(final SpriteLoader.Preparations blockAtlas, final SpriteLoader.Preparations itemAtlas, ModelBakery bakery, Object2IntMap<BlockState> modelGroups, EntityModelSet entityModelSet, SpecialBlockModelRenderer specialBlockModelRenderer, Executor taskExecutor) {
        final Multimap missingMaterials = Multimaps.synchronizedMultimap((Multimap)HashMultimap.create());
        final Multimap missingReferences = Multimaps.synchronizedMultimap((Multimap)HashMultimap.create());
        return bakery.bakeModels(new SpriteGetter(){
            private final TextureAtlasSprite blockMissing;
            private final TextureAtlasSprite itemMissing;
            {
                this.blockMissing = blockAtlas.missing();
                this.itemMissing = itemAtlas.missing();
            }

            @Override
            public TextureAtlasSprite get(Material material, ModelDebugName name) {
                TextureAtlasSprite result;
                Identifier atlasId = material.atlasLocation();
                boolean itemOrBlock = atlasId.equals(BLOCK_OR_ITEM);
                boolean onlyItem = atlasId.equals(TextureAtlas.LOCATION_ITEMS);
                boolean onlyBlock = atlasId.equals(TextureAtlas.LOCATION_BLOCKS);
                if ((itemOrBlock || onlyItem) && (result = itemAtlas.getSprite(material.texture())) != null) {
                    return result;
                }
                if ((itemOrBlock || onlyBlock) && (result = blockAtlas.getSprite(material.texture())) != null) {
                    return result;
                }
                missingMaterials.put((Object)name.debugName(), (Object)material);
                return onlyItem ? this.itemMissing : this.blockMissing;
            }

            @Override
            public TextureAtlasSprite reportMissingReference(String reference, ModelDebugName responsibleModel) {
                missingReferences.put((Object)responsibleModel.debugName(), (Object)reference);
                return this.blockMissing;
            }
        }, taskExecutor).thenApply(bakingResult -> {
            missingMaterials.asMap().forEach((location, materials) -> LOGGER.warn("Missing textures in model {}:\n{}", location, (Object)materials.stream().sorted(Material.COMPARATOR).map(m -> "    " + String.valueOf(m.atlasLocation()) + ":" + String.valueOf(m.texture())).collect(Collectors.joining("\n"))));
            missingReferences.asMap().forEach((location, references) -> LOGGER.warn("Missing texture references in model {}:\n{}", location, (Object)references.stream().sorted().map(reference -> "    " + reference).collect(Collectors.joining("\n"))));
            Map<BlockState, BlockStateModel> modelByStateCache = ModelManager.createBlockStateToModelDispatch(bakingResult.blockStateModels(), bakingResult.missingModels().block());
            return new ReloadState((ModelBakery.BakingResult)bakingResult, modelGroups, modelByStateCache, entityModelSet, specialBlockModelRenderer);
        });
    }

    private static Map<BlockState, BlockStateModel> createBlockStateToModelDispatch(Map<BlockState, BlockStateModel> bakedModels, BlockStateModel missingModel) {
        try (Zone ignored = Profiler.get().zone("block state dispatch");){
            IdentityHashMap<BlockState, BlockStateModel> modelByStateCache = new IdentityHashMap<BlockState, BlockStateModel>(bakedModels);
            for (Block block : BuiltInRegistries.BLOCK) {
                block.getStateDefinition().getPossibleStates().forEach(state -> {
                    if (bakedModels.putIfAbsent((BlockState)state, missingModel) == null) {
                        LOGGER.warn("Missing model for variant: '{}'", state);
                    }
                });
            }
            IdentityHashMap<BlockState, BlockStateModel> identityHashMap = modelByStateCache;
            return identityHashMap;
        }
    }

    private static Object2IntMap<BlockState> buildModelGroups(BlockColors blockColors, BlockStateModelLoader.LoadedModels blockStateModels) {
        try (Zone ignored = Profiler.get().zone("block groups");){
            Object2IntMap<BlockState> object2IntMap = ModelGroupCollector.build(blockColors, blockStateModels);
            return object2IntMap;
        }
    }

    private void apply(ReloadState preparations) {
        ModelBakery.BakingResult bakedModels = preparations.bakedModels;
        this.bakedItemStackModels = bakedModels.itemStackModels();
        this.itemProperties = bakedModels.itemProperties();
        this.modelGroups = preparations.modelGroups;
        this.missingModels = bakedModels.missingModels();
        this.blockModelShaper.replaceCache(preparations.modelCache);
        this.specialBlockModelRenderer = preparations.specialBlockModelRenderer;
        this.entityModelSet = preparations.entityModelSet;
    }

    public boolean requiresRender(BlockState oldState, BlockState newState) {
        int newModelGroup;
        if (oldState == newState) {
            return false;
        }
        int oldModelGroup = this.modelGroups.getInt((Object)oldState);
        if (oldModelGroup != -1 && oldModelGroup == (newModelGroup = this.modelGroups.getInt((Object)newState))) {
            FluidState newFluidState;
            FluidState oldFluidState = oldState.getFluidState();
            return oldFluidState != (newFluidState = newState.getFluidState());
        }
        return true;
    }

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

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

    private /* synthetic */ CompletionStage lambda$reload$4(CompletableFuture pendingBlockAtlasSprites, CompletableFuture pendingItemAtlasSprites, CompletableFuture modelDiscovery, CompletableFuture modelGroups, CompletableFuture modelCache, CompletableFuture entityModelSet, CompletableFuture blockStateModels, CompletableFuture itemStackModels, CompletableFuture specialBlockModelRenderer, Executor taskExecutor, Void unused) {
        SpriteLoader.Preparations blockAtlasSprites = (SpriteLoader.Preparations)pendingBlockAtlasSprites.join();
        SpriteLoader.Preparations itemAtlasSprites = (SpriteLoader.Preparations)pendingItemAtlasSprites.join();
        ResolvedModels resolvedModels = (ResolvedModels)modelDiscovery.join();
        Object2IntMap groups = (Object2IntMap)modelGroups.join();
        Sets.SetView unreferencedModels = Sets.difference(((Map)modelCache.join()).keySet(), resolvedModels.models.keySet());
        if (!unreferencedModels.isEmpty()) {
            LOGGER.debug("Unreferenced models: \n{}", (Object)unreferencedModels.stream().sorted().map(modelId -> "\t" + String.valueOf(modelId) + "\n").collect(Collectors.joining()));
        }
        ModelBakery bakery = new ModelBakery((EntityModelSet)entityModelSet.join(), this.atlasManager, this.playerSkinRenderCache, ((BlockStateModelLoader.LoadedModels)blockStateModels.join()).models(), ((ClientItemInfoLoader.LoadedClientInfos)itemStackModels.join()).contents(), resolvedModels.models(), resolvedModels.missing());
        return ModelManager.loadModels(blockAtlasSprites, itemAtlasSprites, bakery, (Object2IntMap<BlockState>)groups, (EntityModelSet)entityModelSet.join(), (SpecialBlockModelRenderer)specialBlockModelRenderer.join(), taskExecutor);
    }

    private record ResolvedModels(ResolvedModel missing, Map<Identifier, ResolvedModel> models) {
    }

    private record ReloadState(ModelBakery.BakingResult bakedModels, Object2IntMap<BlockState> modelGroups, Map<BlockState, BlockStateModel> modelCache, EntityModelSet entityModelSet, SpecialBlockModelRenderer specialBlockModelRenderer) {
    }
}