package net.minecraft.client.renderer;

import com.mojang.blaze3d.ProjectionType;
import com.mojang.blaze3d.framegraph.FrameGraphBuilder;
import com.mojang.blaze3d.framegraph.FramePass;
import com.mojang.blaze3d.pipeline.RenderTarget;
import com.mojang.blaze3d.resource.ResourceHandle;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.blaze3d.vertex.BufferBuilder;
import com.mojang.blaze3d.vertex.BufferUploader;
import com.mojang.blaze3d.vertex.DefaultVertexFormat;
import com.mojang.blaze3d.vertex.Tesselator;
import com.mojang.blaze3d.vertex.VertexFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.PostChainConfig.Uniform;
import net.minecraft.client.renderer.texture.AbstractTexture;
import net.minecraft.resources.ResourceLocation;
import org.joml.Matrix4f;

@Environment(EnvType.CLIENT)
public class PostPass {
	private final String name;
	private final CompiledShaderProgram shader;
	private final ResourceLocation outputTargetId;
	private final List<Uniform> uniforms;
	private final List<PostPass.Input> inputs = new ArrayList();

	public PostPass(String name, CompiledShaderProgram shader, ResourceLocation outputTargetId, List<Uniform> uniforms) {
		this.name = name;
		this.shader = shader;
		this.outputTargetId = outputTargetId;
		this.uniforms = uniforms;
	}

	public void addInput(PostPass.Input input) {
		this.inputs.add(input);
	}

	public void addToFrame(FrameGraphBuilder frameGraphBuilder, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets, Matrix4f projectionMatrix) {
		FramePass framePass = frameGraphBuilder.addPass(this.name);

		for (PostPass.Input input : this.inputs) {
			input.addToPass(framePass, targets);
		}

		ResourceHandle<RenderTarget> resourceHandle = (ResourceHandle<RenderTarget>)targets.computeIfPresent(
			this.outputTargetId, (resourceLocation, resourceHandlex) -> framePass.readsAndWrites(resourceHandlex)
		);
		if (resourceHandle == null) {
			throw new IllegalStateException("Missing handle for target " + this.outputTargetId);
		} else {
			framePass.executes(() -> {
				RenderTarget renderTarget = resourceHandle.get();
				RenderSystem.viewport(0, 0, renderTarget.width, renderTarget.height);

				for (PostPass.Input inputx : this.inputs) {
					inputx.bindTo(this.shader, targets);
				}

				this.shader.safeGetUniform("OutSize").set((float)renderTarget.width, (float)renderTarget.height);

				for (Uniform uniform : this.uniforms) {
					com.mojang.blaze3d.shaders.Uniform uniform2 = this.shader.getUniform(uniform.name());
					if (uniform2 != null) {
						uniform2.setFromConfig(uniform.values(), uniform.values().size());
					}
				}

				renderTarget.setClearColor(0.0F, 0.0F, 0.0F, 0.0F);
				renderTarget.clear();
				renderTarget.bindWrite(false);
				RenderSystem.depthFunc(519);
				RenderSystem.setShader(this.shader);
				RenderSystem.backupProjectionMatrix();
				RenderSystem.setProjectionMatrix(projectionMatrix, ProjectionType.ORTHOGRAPHIC);
				BufferBuilder bufferBuilder = Tesselator.getInstance().begin(VertexFormat.Mode.QUADS, DefaultVertexFormat.POSITION);
				bufferBuilder.addVertex(0.0F, 0.0F, 500.0F);
				bufferBuilder.addVertex(renderTarget.width, 0.0F, 500.0F);
				bufferBuilder.addVertex(renderTarget.width, renderTarget.height, 500.0F);
				bufferBuilder.addVertex(0.0F, renderTarget.height, 500.0F);
				BufferUploader.drawWithShader(bufferBuilder.buildOrThrow());
				RenderSystem.depthFunc(515);
				RenderSystem.restoreProjectionMatrix();
				renderTarget.unbindWrite();

				for (PostPass.Input input2 : this.inputs) {
					input2.cleanup(targets);
				}

				this.restoreDefaultUniforms();
			});
		}
	}

	private void restoreDefaultUniforms() {
		for (Uniform uniform : this.uniforms) {
			String string = uniform.name();
			com.mojang.blaze3d.shaders.Uniform uniform2 = this.shader.getUniform(string);
			net.minecraft.client.renderer.ShaderProgramConfig.Uniform uniform3 = this.shader.getUniformConfig(string);
			if (uniform2 != null && uniform3 != null && !uniform.values().equals(uniform3.values())) {
				uniform2.setFromConfig(uniform3);
			}
		}
	}

	public CompiledShaderProgram getShader() {
		return this.shader;
	}

	@Environment(EnvType.CLIENT)
	public interface Input {
		void addToPass(FramePass pass, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets);

		void bindTo(CompiledShaderProgram shaderProgram, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets);

		default void cleanup(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
		}
	}

	@Environment(EnvType.CLIENT)
	public record TargetInput(String samplerName, ResourceLocation targetId, boolean depthBuffer, boolean bilinear) implements PostPass.Input {
		private ResourceHandle<RenderTarget> getHandle(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			ResourceHandle<RenderTarget> resourceHandle = (ResourceHandle<RenderTarget>)targets.get(this.targetId);
			if (resourceHandle == null) {
				throw new IllegalStateException("Missing handle for target " + this.targetId);
			} else {
				return resourceHandle;
			}
		}

		@Override
		public void addToPass(FramePass pass, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			pass.reads(this.getHandle(targets));
		}

		@Override
		public void bindTo(CompiledShaderProgram shaderProgram, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			ResourceHandle<RenderTarget> resourceHandle = this.getHandle(targets);
			RenderTarget renderTarget = resourceHandle.get();
			renderTarget.setFilterMode(this.bilinear ? 9729 : 9728);
			shaderProgram.bindSampler(this.samplerName + "Sampler", this.depthBuffer ? renderTarget.getDepthTextureId() : renderTarget.getColorTextureId());
			shaderProgram.safeGetUniform(this.samplerName + "Size").set((float)renderTarget.width, (float)renderTarget.height);
		}

		@Override
		public void cleanup(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			if (this.bilinear) {
				this.getHandle(targets).get().setFilterMode(9728);
			}
		}
	}

	@Environment(EnvType.CLIENT)
	public record TextureInput(String samplerName, AbstractTexture texture, int width, int height) implements PostPass.Input {
		@Override
		public void addToPass(FramePass pass, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
		}

		@Override
		public void bindTo(CompiledShaderProgram shaderProgram, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			shaderProgram.bindSampler(this.samplerName + "Sampler", this.texture.getId());
			shaderProgram.safeGetUniform(this.samplerName + "Size").set((float)this.width, (float)this.height);
		}
	}
}