package net.minecraft.client.renderer;

import com.mojang.blaze3d.ProjectionType;
import com.mojang.blaze3d.buffers.GpuBuffer;
import com.mojang.blaze3d.framegraph.FrameGraphBuilder;
import com.mojang.blaze3d.framegraph.FramePass;
import com.mojang.blaze3d.pipeline.RenderPipeline;
import com.mojang.blaze3d.pipeline.RenderTarget;
import com.mojang.blaze3d.resource.ResourceHandle;
import com.mojang.blaze3d.systems.RenderPass;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.blaze3d.textures.FilterMode;
import com.mojang.blaze3d.textures.GpuTexture;
import com.mojang.blaze3d.vertex.VertexFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.OptionalDouble;
import java.util.OptionalInt;
import java.util.function.Consumer;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.texture.AbstractTexture;
import net.minecraft.resources.ResourceLocation;
import org.jetbrains.annotations.Nullable;
import org.joml.Matrix4f;

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

	public PostPass(RenderPipeline pipeline, ResourceLocation outputTargetId, List<PostChainConfig.Uniform> uniforms) {
		this.pipeline = pipeline;
		this.name = pipeline.getLocation().toString();
		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,
		@Nullable Consumer<RenderPass> uniformSetter
	) {
		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.backupProjectionMatrix();
					RenderSystem.setProjectionMatrix(projectionMatrix, ProjectionType.ORTHOGRAPHIC);
					GpuBuffer gpuBuffer = RenderSystem.getQuadVertexBuffer();
					RenderSystem.AutoStorageIndexBuffer autoStorageIndexBuffer = RenderSystem.getSequentialBuffer(VertexFormat.Mode.QUADS);
					GpuBuffer gpuBuffer2 = autoStorageIndexBuffer.getBuffer(6);

					try (RenderPass renderPass = RenderSystem.getDevice()
							.createCommandEncoder()
							.createRenderPass(
								renderTarget.getColorTexture(), OptionalInt.empty(), renderTarget.useDepth ? renderTarget.getDepthTexture() : null, OptionalDouble.empty()
							)) {
						renderPass.setPipeline(this.pipeline);
						renderPass.setUniform("OutSize", (float)renderTarget.width, (float)renderTarget.height);
						renderPass.setVertexBuffer(0, gpuBuffer);
						renderPass.setIndexBuffer(gpuBuffer2, autoStorageIndexBuffer.type());

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

						if (uniformSetter != null) {
							uniformSetter.accept(renderPass);
						}

						for (PostChainConfig.Uniform uniform : this.uniforms) {
							uniform.setOnRenderPass(renderPass);
						}

						renderPass.drawIndexed(0, 6);
					}

					RenderSystem.restoreProjectionMatrix();

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

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

		void bindTo(RenderPass renderPass, 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(RenderPass renderPass, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			ResourceHandle<RenderTarget> resourceHandle = this.getHandle(targets);
			RenderTarget renderTarget = resourceHandle.get();
			renderTarget.setFilterMode(this.bilinear ? FilterMode.LINEAR : FilterMode.NEAREST);
			GpuTexture gpuTexture = this.depthBuffer ? renderTarget.getDepthTexture() : renderTarget.getColorTexture();
			if (gpuTexture == null) {
				throw new IllegalStateException("Missing " + (this.depthBuffer ? "depth" : "color") + "texture for target " + this.targetId);
			} else {
				renderPass.bindSampler(this.samplerName + "Sampler", gpuTexture);
				renderPass.setUniform(this.samplerName + "Size", (float)renderTarget.width, (float)renderTarget.height);
			}
		}

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

	@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(RenderPass renderPass, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			renderPass.bindSampler(this.samplerName + "Sampler", this.texture.getTexture());
			renderPass.setUniform(this.samplerName + "Size", (float)this.width, (float)this.height);
		}
	}
}