package net.minecraft.client.renderer;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Sets;
import com.google.common.collect.ImmutableList.Builder;
import com.mojang.blaze3d.buffers.GpuBufferSlice;
import com.mojang.blaze3d.framegraph.FrameGraphBuilder;
import com.mojang.blaze3d.pipeline.RenderPipeline;
import com.mojang.blaze3d.pipeline.RenderTarget;
import com.mojang.blaze3d.resource.GraphicsResourceAllocator;
import com.mojang.blaze3d.resource.RenderTargetDescriptor;
import com.mojang.blaze3d.resource.ResourceHandle;
import com.mojang.blaze3d.shaders.UniformType;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.texture.AbstractTexture;
import net.minecraft.client.renderer.texture.TextureManager;
import net.minecraft.resources.ResourceLocation;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class PostChain implements AutoCloseable {
	public static final ResourceLocation MAIN_TARGET_ID = ResourceLocation.withDefaultNamespace("main");
	private final List<PostPass> passes;
	private final Map<ResourceLocation, PostChainConfig.InternalTarget> internalTargets;
	private final Set<ResourceLocation> externalTargets;
	private final Map<ResourceLocation, RenderTarget> persistentTargets = new HashMap();
	private final CachedOrthoProjectionMatrixBuffer projectionMatrixBuffer;

	private PostChain(
		List<PostPass> passes,
		Map<ResourceLocation, PostChainConfig.InternalTarget> internalTargets,
		Set<ResourceLocation> externalTargets,
		CachedOrthoProjectionMatrixBuffer projectionMatrixBuffer
	) {
		this.passes = passes;
		this.internalTargets = internalTargets;
		this.externalTargets = externalTargets;
		this.projectionMatrixBuffer = projectionMatrixBuffer;
	}

	public static PostChain load(
		PostChainConfig config,
		TextureManager textureManager,
		Set<ResourceLocation> externalTargets,
		ResourceLocation name,
		CachedOrthoProjectionMatrixBuffer projectionMatrixBuffer
	) throws ShaderManager.CompilationException {
		Stream<ResourceLocation> stream = config.passes().stream().flatMap(PostChainConfig.Pass::referencedTargets);
		Set<ResourceLocation> set = (Set<ResourceLocation>)stream.filter(resourceLocation -> !config.internalTargets().containsKey(resourceLocation))
			.collect(Collectors.toSet());
		Set<ResourceLocation> set2 = Sets.<ResourceLocation>difference(set, externalTargets);
		if (!set2.isEmpty()) {
			throw new ShaderManager.CompilationException("Referenced external targets are not available in this context: " + set2);
		} else {
			Builder<PostPass> builder = ImmutableList.builder();

			for (int i = 0; i < config.passes().size(); i++) {
				PostChainConfig.Pass pass = (PostChainConfig.Pass)config.passes().get(i);
				builder.add(createPass(textureManager, pass, name.withSuffix("/" + i)));
			}

			return new PostChain(builder.build(), config.internalTargets(), set, projectionMatrixBuffer);
		}
	}

	private static PostPass createPass(TextureManager textureManager, PostChainConfig.Pass pass, ResourceLocation location) throws ShaderManager.CompilationException {
		RenderPipeline.Builder builder = RenderPipeline.builder(RenderPipelines.POST_PROCESSING_SNIPPET)
			.withFragmentShader(pass.fragmentShaderId())
			.withVertexShader(pass.vertexShaderId())
			.withLocation(location);

		for (PostChainConfig.Input input : pass.inputs()) {
			builder.withSampler(input.samplerName() + "Sampler");
		}

		builder.withUniform("SamplerInfo", UniformType.UNIFORM_BUFFER);

		for (String string : pass.uniforms().keySet()) {
			builder.withUniform(string, UniformType.UNIFORM_BUFFER);
		}

		RenderPipeline renderPipeline = builder.build();
		List<PostPass.Input> list = new ArrayList();

		for (PostChainConfig.Input input2 : pass.inputs()) {
			switch (input2) {
				case PostChainConfig.TextureInput(String var35, ResourceLocation var36, int var37, int var38, boolean var39):
					AbstractTexture abstractTexture = textureManager.getTexture(var36.withPath((UnaryOperator<String>)(string -> "textures/effect/" + string + ".png")));
					abstractTexture.setFilter(var39, false);
					list.add(new PostPass.TextureInput(var35, abstractTexture, var37, var38));
					break;
				case PostChainConfig.TargetInput(String var21, ResourceLocation var41, boolean var42, boolean var43):
					list.add(new PostPass.TargetInput(var21, var41, var42, var43));
					break;
				default:
					throw new MatchException(null, null);
			}
		}

		return new PostPass(renderPipeline, pass.outputTarget(), pass.uniforms(), list);
	}

	public void addToFrame(FrameGraphBuilder frameGraphBuilder, int width, int height, PostChain.TargetBundle targetBundle) {
		GpuBufferSlice gpuBufferSlice = this.projectionMatrixBuffer.getBuffer(width, height);
		Map<ResourceLocation, ResourceHandle<RenderTarget>> map = new HashMap(this.internalTargets.size() + this.externalTargets.size());

		for (ResourceLocation resourceLocation : this.externalTargets) {
			map.put(resourceLocation, targetBundle.getOrThrow(resourceLocation));
		}

		for (Entry<ResourceLocation, PostChainConfig.InternalTarget> entry : this.internalTargets.entrySet()) {
			ResourceLocation resourceLocation2 = (ResourceLocation)entry.getKey();
			PostChainConfig.InternalTarget internalTarget = (PostChainConfig.InternalTarget)entry.getValue();
			RenderTargetDescriptor renderTargetDescriptor = new RenderTargetDescriptor(
				(Integer)internalTarget.width().orElse(width), (Integer)internalTarget.height().orElse(height), true, internalTarget.clearColor()
			);
			if (internalTarget.persistent()) {
				RenderTarget renderTarget = this.getOrCreatePersistentTarget(resourceLocation2, renderTargetDescriptor);
				map.put(resourceLocation2, frameGraphBuilder.importExternal(resourceLocation2.toString(), renderTarget));
			} else {
				map.put(resourceLocation2, frameGraphBuilder.createInternal(resourceLocation2.toString(), renderTargetDescriptor));
			}
		}

		for (PostPass postPass : this.passes) {
			postPass.addToFrame(frameGraphBuilder, map, gpuBufferSlice);
		}

		for (ResourceLocation resourceLocation : this.externalTargets) {
			targetBundle.replace(resourceLocation, (ResourceHandle<RenderTarget>)map.get(resourceLocation));
		}
	}

	@Deprecated
	public void process(RenderTarget target, GraphicsResourceAllocator graphicsResourceAllocator) {
		FrameGraphBuilder frameGraphBuilder = new FrameGraphBuilder();
		PostChain.TargetBundle targetBundle = PostChain.TargetBundle.of(MAIN_TARGET_ID, frameGraphBuilder.importExternal("main", target));
		this.addToFrame(frameGraphBuilder, target.width, target.height, targetBundle);
		frameGraphBuilder.execute(graphicsResourceAllocator);
	}

	private RenderTarget getOrCreatePersistentTarget(ResourceLocation name, RenderTargetDescriptor descriptor) {
		RenderTarget renderTarget = (RenderTarget)this.persistentTargets.get(name);
		if (renderTarget == null || renderTarget.width != descriptor.width() || renderTarget.height != descriptor.height()) {
			if (renderTarget != null) {
				renderTarget.destroyBuffers();
			}

			renderTarget = descriptor.allocate();
			descriptor.prepare(renderTarget);
			this.persistentTargets.put(name, renderTarget);
		}

		return renderTarget;
	}

	public void close() {
		this.persistentTargets.values().forEach(RenderTarget::destroyBuffers);
		this.persistentTargets.clear();

		for (PostPass postPass : this.passes) {
			postPass.close();
		}
	}

	@Environment(EnvType.CLIENT)
	public interface TargetBundle {
		static PostChain.TargetBundle of(ResourceLocation id, ResourceHandle<RenderTarget> handle) {
			return new PostChain.TargetBundle() {
				private ResourceHandle<RenderTarget> handle = handle;

				@Override
				public void replace(ResourceLocation id, ResourceHandle<RenderTarget> handle) {
					if (id.equals(id)) {
						this.handle = handle;
					} else {
						throw new IllegalArgumentException("No target with id " + id);
					}
				}

				@Nullable
				@Override
				public ResourceHandle<RenderTarget> get(ResourceLocation id) {
					return id.equals(id) ? this.handle : null;
				}
			};
		}

		void replace(ResourceLocation id, ResourceHandle<RenderTarget> handle);

		@Nullable
		ResourceHandle<RenderTarget> get(ResourceLocation id);

		default ResourceHandle<RenderTarget> getOrThrow(ResourceLocation id) {
			ResourceHandle<RenderTarget> resourceHandle = this.get(id);
			if (resourceHandle == null) {
				throw new IllegalArgumentException("Missing target with id " + id);
			} else {
				return resourceHandle;
			}
		}
	}
}