package net.minecraft.client.renderer;

import com.mojang.blaze3d.ProjectionType;
import com.mojang.blaze3d.buffers.GpuBuffer;
import com.mojang.blaze3d.buffers.GpuBufferSlice;
import com.mojang.blaze3d.buffers.Std140Builder;
import com.mojang.blaze3d.buffers.Std140SizeCalculator;
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.CommandEncoder;
import com.mojang.blaze3d.systems.RenderPass;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.blaze3d.textures.FilterMode;
import com.mojang.blaze3d.textures.GpuTextureView;
import com.mojang.blaze3d.vertex.VertexFormat;
import com.mojang.datafixers.util.Pair;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.OptionalDouble;
import java.util.OptionalInt;
import java.util.Map.Entry;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.texture.AbstractTexture;
import net.minecraft.resources.ResourceLocation;
import org.lwjgl.system.MemoryStack;

@Environment(EnvType.CLIENT)
public class PostPass implements AutoCloseable {
	private static final int UBO_SIZE_PER_SAMPLER = new Std140SizeCalculator().putVec2().get();
	private final String name;
	private final RenderPipeline pipeline;
	private final ResourceLocation outputTargetId;
	private final Map<String, GpuBuffer> customUniforms = new HashMap();
	private final MappableRingBuffer infoUbo;
	private final List<PostPass.Input> inputs;

	public PostPass(RenderPipeline pipeline, ResourceLocation outputTargetId, Map<String, List<UniformValue>> uniforms, List<PostPass.Input> inputs) {
		this.pipeline = pipeline;
		this.name = pipeline.getLocation().toString();
		this.outputTargetId = outputTargetId;
		this.inputs = inputs;

		for (Entry<String, List<UniformValue>> entry : uniforms.entrySet()) {
			List<UniformValue> list = (List<UniformValue>)entry.getValue();
			if (!list.isEmpty()) {
				Std140SizeCalculator std140SizeCalculator = new Std140SizeCalculator();

				for (UniformValue uniformValue : list) {
					uniformValue.addSize(std140SizeCalculator);
				}

				int i = std140SizeCalculator.get();

				try (MemoryStack memoryStack = MemoryStack.stackPush()) {
					Std140Builder std140Builder = Std140Builder.onStack(memoryStack, i);

					for (UniformValue uniformValue2 : list) {
						uniformValue2.writeTo(std140Builder);
					}

					this.customUniforms
						.put((String)entry.getKey(), RenderSystem.getDevice().createBuffer(() -> this.name + " / " + (String)entry.getKey(), 128, std140Builder.get()));
				}
			}
		}

		this.infoUbo = new MappableRingBuffer(() -> this.name + " SamplerInfo", 130, (inputs.size() + 1) * UBO_SIZE_PER_SAMPLER);
	}

	public void addToFrame(FrameGraphBuilder frameGraphBuilder, Map<ResourceLocation, ResourceHandle<RenderTarget>> targets, GpuBufferSlice projectionMatrixBuffer) {
		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(projectionMatrixBuffer, ProjectionType.ORTHOGRAPHIC);
					CommandEncoder commandEncoder = RenderSystem.getDevice().createCommandEncoder();
					List<Pair<String, GpuTextureView>> list = this.inputs.stream().map(inputxx -> Pair.of(inputxx.samplerName(), inputxx.texture(targets))).toList();

					try (GpuBuffer.MappedView mappedView = commandEncoder.mapBuffer(this.infoUbo.currentBuffer(), false, true)) {
						Std140Builder std140Builder = Std140Builder.intoBuffer(mappedView.data());
						std140Builder.putVec2(renderTarget.width, renderTarget.height);

						for (Pair<String, GpuTextureView> pair : list) {
							std140Builder.putVec2(pair.getSecond().getWidth(0), pair.getSecond().getHeight(0));
						}
					}

					GpuBuffer gpuBuffer = RenderSystem.getQuadVertexBuffer();
					RenderSystem.AutoStorageIndexBuffer autoStorageIndexBuffer = RenderSystem.getSequentialBuffer(VertexFormat.Mode.QUADS);
					GpuBuffer gpuBuffer2 = autoStorageIndexBuffer.getBuffer(6);

					try (RenderPass renderPass = commandEncoder.createRenderPass(
							() -> "Post pass " + this.name,
							renderTarget.getColorTextureView(),
							OptionalInt.empty(),
							renderTarget.useDepth ? renderTarget.getDepthTextureView() : null,
							OptionalDouble.empty()
						)) {
						renderPass.setPipeline(this.pipeline);
						RenderSystem.bindDefaultUniforms(renderPass);
						renderPass.setUniform("SamplerInfo", this.infoUbo.currentBuffer());

						for (Entry<String, GpuBuffer> entry : this.customUniforms.entrySet()) {
							renderPass.setUniform((String)entry.getKey(), (GpuBuffer)entry.getValue());
						}

						renderPass.setVertexBuffer(0, gpuBuffer);
						renderPass.setIndexBuffer(gpuBuffer2, autoStorageIndexBuffer.type());

						for (Pair<String, GpuTextureView> pair2 : list) {
							renderPass.bindSampler(pair2.getFirst() + "Sampler", pair2.getSecond());
						}

						renderPass.drawIndexed(0, 0, 6, 1);
					}

					this.infoUbo.rotate();
					RenderSystem.restoreProjectionMatrix();

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

	public void close() {
		for (GpuBuffer gpuBuffer : this.customUniforms.values()) {
			gpuBuffer.close();
		}

		this.infoUbo.close();
	}

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

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

		GpuTextureView texture(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets);

		String samplerName();
	}

	@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 cleanup(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			if (this.bilinear) {
				this.getHandle(targets).get().setFilterMode(FilterMode.NEAREST);
			}
		}

		@Override
		public GpuTextureView texture(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			ResourceHandle<RenderTarget> resourceHandle = this.getHandle(targets);
			RenderTarget renderTarget = resourceHandle.get();
			renderTarget.setFilterMode(this.bilinear ? FilterMode.LINEAR : FilterMode.NEAREST);
			GpuTextureView gpuTextureView = this.depthBuffer ? renderTarget.getDepthTextureView() : renderTarget.getColorTextureView();
			if (gpuTextureView == null) {
				throw new IllegalStateException("Missing " + (this.depthBuffer ? "depth" : "color") + "texture for target " + this.targetId);
			} else {
				return gpuTextureView;
			}
		}
	}

	@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 GpuTextureView texture(Map<ResourceLocation, ResourceHandle<RenderTarget>> targets) {
			return this.texture.getTextureView();
		}
	}
}