package net.minecraft.client.resources.model.sprite;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.mojang.logging.LogUtils;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.stream.Collectors;
import net.minecraft.client.renderer.texture.MissingTextureAtlasSprite;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.renderer.texture.SpriteLoader.Preparations;
import net.minecraft.client.resources.model.ModelDebugName;
import net.minecraft.client.resources.model.sprite.Material.Baked;
import net.minecraft.resources.Identifier;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;

public abstract class MaterialBaker {
	private static final Logger LOGGER = LogUtils.getLogger();
	private final Baked missingSprite;
	private final Baked missingSpriteForceTranslucent;
	private final Multimap<String, Identifier> missingSprites = Multimaps.synchronizedMultimap(HashMultimap.create());
	private final Multimap<String, String> missingReferences = Multimaps.synchronizedMultimap(HashMultimap.create());
	private final Map<Material, Baked> bakedMaterials = new ConcurrentHashMap();
	private final Function<Material, Baked> bakerFunction = this::bake;

	public MaterialBaker(final TextureAtlasSprite missingSprite) {
		this.missingSprite = new Baked(missingSprite, false);
		this.missingSpriteForceTranslucent = new Baked(missingSprite, true);
	}

	public Baked replacementForMissingMaterial(final Material material) {
		return material.forceTranslucent() ? this.missingSpriteForceTranslucent : this.missingSprite;
	}

	public Baked get(final Material material, final ModelDebugName name) {
		if (material.sprite().equals(MissingTextureAtlasSprite.getLocation())) {
			return this.replacementForMissingMaterial(material);
		} else {
			Baked baked = (Baked)this.bakedMaterials.computeIfAbsent(material, this.bakerFunction);
			if (baked == null) {
				this.missingSprites.put(name.debugName(), material.sprite());
				return this.replacementForMissingMaterial(material);
			} else {
				return baked;
			}
		}
	}

	@Nullable
	protected abstract Baked bake(Material material);

	@Nullable
	protected static Baked bakeForAtlas(final Material material, final Preparations atlas) {
		TextureAtlasSprite sprite = atlas.getSprite(material.sprite());
		return sprite != null ? new Baked(sprite, material.forceTranslucent()) : null;
	}

	public Baked resolveSlot(final TextureSlots slots, final String id, final ModelDebugName name) {
		Material resolvedMaterial = slots.getMaterial(id);
		return resolvedMaterial != null ? this.get(resolvedMaterial, name) : this.reportMissingReference(id, name);
	}

	public Baked reportMissingReference(final String reference, final ModelDebugName responsibleModel) {
		this.missingReferences.put(responsibleModel.debugName(), reference);
		return this.missingSprite;
	}

	public void logMissingTextures() {
		this.missingSprites
			.asMap()
			.forEach(
				(location, sprites) -> LOGGER.warn(
					"Missing textures in model {}:\n{}", location, sprites.stream().sorted().map(sprite -> "    " + sprite).collect(Collectors.joining("\n"))
				)
			);
		this.missingReferences
			.asMap()
			.forEach(
				(location, references) -> LOGGER.warn(
					"Missing texture references in model {}:\n{}",
					location,
					references.stream().sorted().map(reference -> "    " + reference).collect(Collectors.joining("\n"))
				)
			);
	}
}