package net.minecraft.client.resources.model;

import com.google.common.collect.Sets;
import com.mojang.logging.LogUtils;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.block.model.ItemModelGenerator;
import net.minecraft.resources.ResourceLocation;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class ModelDiscovery {
	static final Logger LOGGER = LogUtils.getLogger();
	private final Map<ResourceLocation, UnbakedModel> inputModels;
	final UnbakedModel missingModel;
	private final List<ResolvableModel> topModels = new ArrayList();
	private final Map<ResourceLocation, UnbakedModel> referencedModels = new HashMap();

	public ModelDiscovery(Map<ResourceLocation, UnbakedModel> inputModels, UnbakedModel missingModel) {
		this.inputModels = inputModels;
		this.missingModel = missingModel;
		this.referencedModels.put(MissingBlockModel.LOCATION, missingModel);
	}

	public void registerSpecialModels() {
		this.referencedModels.put(ItemModelGenerator.GENERATED_ITEM_MODEL_ID, new ItemModelGenerator());
	}

	public void addRoot(ResolvableModel model) {
		this.topModels.add(model);
	}

	public void discoverDependencies() {
		this.topModels.forEach(resolvableModel -> resolvableModel.resolveDependencies(new ModelDiscovery.ResolverImpl()));
	}

	public Map<ResourceLocation, UnbakedModel> getReferencedModels() {
		return this.referencedModels;
	}

	public Set<ResourceLocation> getUnreferencedModels() {
		return Sets.<ResourceLocation>difference(this.inputModels.keySet(), this.referencedModels.keySet());
	}

	UnbakedModel getBlockModel(ResourceLocation modelLocation) {
		return (UnbakedModel)this.referencedModels.computeIfAbsent(modelLocation, this::loadBlockModel);
	}

	private UnbakedModel loadBlockModel(ResourceLocation modelLocation) {
		UnbakedModel unbakedModel = (UnbakedModel)this.inputModels.get(modelLocation);
		if (unbakedModel == null) {
			LOGGER.warn("Missing block model: '{}'", modelLocation);
			return this.missingModel;
		} else {
			return unbakedModel;
		}
	}

	@Environment(EnvType.CLIENT)
	class ResolverImpl implements ResolvableModel.Resolver {
		private final List<ResourceLocation> stack = new ArrayList();
		private final Set<ResourceLocation> resolvedModels = new HashSet();

		@Override
		public UnbakedModel resolve(ResourceLocation model) {
			if (this.stack.contains(model)) {
				ModelDiscovery.LOGGER.warn("Detected model loading loop: {}->{}", this.stacktraceToString(), model);
				return ModelDiscovery.this.missingModel;
			} else {
				UnbakedModel unbakedModel = ModelDiscovery.this.getBlockModel(model);
				if (this.resolvedModels.add(model)) {
					this.stack.add(model);
					unbakedModel.resolveDependencies(this);
					this.stack.remove(model);
				}

				return unbakedModel;
			}
		}

		private String stacktraceToString() {
			return (String)this.stack.stream().map(ResourceLocation::toString).collect(Collectors.joining("->"));
		}
	}
}