package net.minecraft.client.resources;

import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
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.renderer.texture.TextureManager;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.packs.metadata.MetadataSectionType;
import net.minecraft.server.packs.resources.PreparableReloadListener;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.util.profiling.Profiler;
import net.minecraft.util.profiling.Zone;

@Environment(EnvType.CLIENT)
public abstract class TextureAtlasHolder implements PreparableReloadListener, AutoCloseable {
	private final TextureAtlas textureAtlas;
	private final ResourceLocation atlasInfoLocation;
	private final Set<MetadataSectionType<?>> metadataSections;

	public TextureAtlasHolder(TextureManager textureManager, ResourceLocation textureAtlasLocation, ResourceLocation atlasInfoLocation) {
		this(textureManager, textureAtlasLocation, atlasInfoLocation, SpriteLoader.DEFAULT_METADATA_SECTIONS);
	}

	public TextureAtlasHolder(
		TextureManager textureManager, ResourceLocation textureAtlasLocation, ResourceLocation atlasInfoLocation, Set<MetadataSectionType<?>> metadataSections
	) {
		this.atlasInfoLocation = atlasInfoLocation;
		this.textureAtlas = new TextureAtlas(textureAtlasLocation);
		textureManager.register(this.textureAtlas.location(), this.textureAtlas);
		this.metadataSections = metadataSections;
	}

	/**
	 * Gets a sprite associated with the passed resource location.
	 */
	protected TextureAtlasSprite getSprite(ResourceLocation location) {
		return this.textureAtlas.getSprite(location);
	}

	@Override
	public final CompletableFuture<Void> reload(
		PreparableReloadListener.PreparationBarrier preparationBarrier, ResourceManager resourceManager, Executor executor, Executor executor2
	) {
		return SpriteLoader.create(this.textureAtlas)
			.loadAndStitch(resourceManager, this.atlasInfoLocation, 0, executor, this.metadataSections)
			.thenCompose(SpriteLoader.Preparations::waitForUpload)
			.thenCompose(preparationBarrier::wait)
			.thenAcceptAsync(this::apply, executor2);
	}

	private void apply(SpriteLoader.Preparations preperations) {
		try (Zone zone = Profiler.get().zone("upload")) {
			this.textureAtlas.upload(preperations);
		}
	}

	public void close() {
		this.textureAtlas.clearTextureData();
	}
}