package com.mojang.blaze3d.systems;

import com.mojang.blaze3d.IndexType;
import com.mojang.blaze3d.buffers.GpuBuffer;
import com.mojang.blaze3d.buffers.GpuBufferSlice;
import com.mojang.blaze3d.pipeline.ColorTargetState;
import com.mojang.blaze3d.pipeline.RenderPipeline;
import com.mojang.blaze3d.systems.RenderPassDescriptor.Attachment;
import com.mojang.blaze3d.textures.GpuSampler;
import com.mojang.blaze3d.textures.GpuTextureView;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.function.BiConsumer;
import java.util.function.Supplier;
import org.joml.Vector4fc;
import org.jspecify.annotations.Nullable;
import org.lwjgl.vulkan.VkDrawIndexedIndirectCommand;
import org.lwjgl.vulkan.VkDrawIndirectCommand;

public class RenderPass implements AutoCloseable {
	public static final int MAX_VERTEX_BUFFERS = 16;
	private final RenderPassBackend backend;
	private final GpuDeviceBackend device;
	private final DeviceFeatures deviceFeatures;
	private final Runnable onFinish;
	@Nullable
	private final RenderPass.RenderArea renderArea;
	private boolean isClosed;
	private int pushedDebugGroups;
	private final List<Attachment<Optional<Vector4fc>>> colorAttachments;

	public RenderPass(
		final RenderPassBackend backend,
		final GpuDeviceBackend device,
		final List<Attachment<Optional<Vector4fc>>> colorAttachments,
		final Runnable onFinish,
		@Nullable final RenderPass.RenderArea renderArea
	) {
		this.backend = backend;
		this.device = device;
		this.deviceFeatures = device.getDeviceInfo().features();
		this.colorAttachments = colorAttachments;
		this.onFinish = onFinish;
		this.renderArea = renderArea;
	}

	public void pushDebugGroup(final Supplier<String> label) {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else {
			this.pushedDebugGroups++;
			this.backend.pushDebugGroup(label);
		}
	}

	public void popDebugGroup() {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else if (this.pushedDebugGroups == 0) {
			throw new IllegalStateException("Can't pop more debug groups than was pushed!");
		} else {
			this.pushedDebugGroups--;
			this.backend.popDebugGroup();
		}
	}

	public void writeTimestamp(final GpuQueryPool pool, final int index) {
		if (index >= 0 && index <= pool.size()) {
			this.backend.writeTimestamp(pool, index);
		} else {
			throw new IllegalStateException("Index " + index + " is out of range for query pool of size " + pool.size());
		}
	}

	public void setPipeline(final RenderPipeline pipeline) {
		ColorTargetState[] colorTargetStates = pipeline.getColorTargetStates();
		if (colorTargetStates.length != this.colorAttachments.size()) {
			throw new IllegalStateException("Render pass color attachment count must match pipeline color target state count.");
		} else {
			for (int i = 0; i < this.colorAttachments.size(); i++) {
				Attachment<Optional<Vector4fc>> attachment = (Attachment<Optional<Vector4fc>>)this.colorAttachments.get(i);
				if (attachment != null) {
					ColorTargetState colorTargetState = colorTargetStates[i];
					if (colorTargetState == null || colorTargetState.format() != attachment.textureView().texture().getFormat()) {
						throw new IllegalStateException("Render pass color attachment " + i + " format doesn't match pipeline format.");
					}
				}
			}

			this.backend.setPipeline(pipeline);
		}
	}

	public void bindTexture(final String name, @Nullable final GpuTextureView textureView, @Nullable final GpuSampler sampler) {
		this.backend.bindTexture(name, textureView, sampler);
	}

	public void setUniform(final String name, final GpuBuffer value) {
		this.backend.setUniform(name, value);
	}

	public void setUniform(final String name, final GpuBufferSlice value) {
		int alignment = this.device.getDeviceInfo().limits().minUniformOffsetAlignment();
		if (value.offset() % alignment > 0L) {
			throw new IllegalArgumentException("Uniform buffer offset must be aligned to " + alignment);
		} else {
			this.backend.setUniform(name, value);
		}
	}

	public void enableScissor(final int x, final int y, final int width, final int height) {
		if (width > 0 && height > 0) {
			if (x >= this.renderArea.x()
				&& y >= this.renderArea.y()
				&& x + width <= this.renderArea.x() + this.renderArea.width()
				&& y + height <= this.renderArea.height()) {
				this.backend.enableScissor(x, y, width, height);
			} else {
				throw new IllegalArgumentException(
					"Scissor at " + x + ", " + y + " with size " + width + "x" + height + " is out of bounds for render area " + this.renderArea
				);
			}
		} else {
			throw new IllegalArgumentException("Scissor size must be >0, was " + width + "x" + height);
		}
	}

	public void disableScissor() {
		this.backend.disableScissor();
	}

	public void setVertexBuffer(final int slot, @Nullable final GpuBufferSlice vertexBuffer) {
		if (slot < 0 || slot >= 16) {
			throw new IllegalArgumentException("Vertex buffer slot is out of range: " + slot);
		} else if (vertexBuffer != null && vertexBuffer.buffer().isClosed()) {
			throw new IllegalStateException("Vertex buffer at slot " + slot + " has been closed!");
		} else if (vertexBuffer != null && (vertexBuffer.buffer().usage() & 32) == 0) {
			throw new IllegalStateException("Vertex buffer at slot " + slot + " doesn't have GpuBuffer.USAGE_VERTEX flag!");
		} else {
			this.backend.setVertexBuffer(slot, vertexBuffer);
		}
	}

	public void setIndexBuffer(final GpuBuffer indexBuffer, final IndexType indexType) {
		this.backend.setIndexBuffer(indexBuffer, indexType);
	}

	public void drawIndexed(final int indexCount, final int instanceCount, final int firstIndex, final int vertexOffset, final int firstInstance) {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else if (firstInstance != 0 && !this.deviceFeatures.nonZeroFirstInstance()) {
			throw new IllegalArgumentException("firstInstance must be zero on when device does not support nonZeroFirstInstance");
		} else {
			this.backend.drawIndexed(indexCount, instanceCount, firstIndex, vertexOffset, firstInstance);
		}
	}

	public void drawIndexedIndirect(final GpuBufferSlice commands, final int drawCount) {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else if (!this.deviceFeatures.drawIndirect()) {
			throw new IllegalArgumentException("device does not support drawIndirect");
		} else if (drawCount > 1 && !this.deviceFeatures.multiDrawIndirect()) {
			throw new IllegalArgumentException("drawCount must be one when device does not support multiDrawIndirect");
		} else if ((commands.buffer().usage() & 512) == 0) {
			throw new IllegalArgumentException("Indirect commands buffer must have GpuBuffer.USAGE_INDIRECT_PARAMETERS flag");
		} else if (commands.length() < (long)drawCount * VkDrawIndexedIndirectCommand.SIZEOF) {
			throw new IllegalArgumentException("Commands buffer is not large enough to hold requested draw count at the given offset");
		} else if (commands.offset() % 4L != 0L) {
			throw new IllegalArgumentException("Commands offset must be multiple of 4");
		} else {
			this.backend.drawIndexedIndirect(commands, drawCount);
		}
	}

	public <T> void drawMultipleIndexed(
		final Collection<RenderPass.Draw<T>> draws,
		@Nullable final GpuBuffer defaultIndexBuffer,
		@Nullable final IndexType defaultIndexType,
		final Collection<String> dynamicUniforms,
		final T uniformArgument
	) {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else {
			this.backend.drawMultipleIndexed(draws, defaultIndexBuffer, defaultIndexType, dynamicUniforms, uniformArgument);
		}
	}

	public void draw(final int vertexCount, final int instanceCount, final int firstVertex, final int firstInstance) {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else if (firstInstance != 0 && !this.deviceFeatures.nonZeroFirstInstance()) {
			throw new IllegalArgumentException("firstInstance must be zero on when device does not support nonZeroFirstInstance");
		} else {
			this.backend.draw(vertexCount, instanceCount, firstVertex, firstInstance);
		}
	}

	public void drawIndirect(final GpuBufferSlice commands, final int drawCount) {
		if (this.isClosed) {
			throw new IllegalStateException("Can't use a closed render pass");
		} else if (!this.deviceFeatures.drawIndirect()) {
			throw new IllegalArgumentException("device does not support drawIndirect");
		} else if (drawCount > 1 && !this.deviceFeatures.multiDrawIndirect()) {
			throw new IllegalArgumentException("drawCount must be one when device does not support multiDrawIndirect");
		} else if ((commands.buffer().usage() & 512) == 0) {
			throw new IllegalArgumentException("Indirect commands buffer must have GpuBuffer.USAGE_INDIRECT_PARAMETERS flag");
		} else if (commands.length() < (long)drawCount * VkDrawIndirectCommand.SIZEOF) {
			throw new IllegalArgumentException("Commands buffer is not large enough to hold requested draw count at the given offset");
		} else if (commands.offset() % 4L != 0L) {
			throw new IllegalArgumentException("Commands offset must be multiple of 4");
		} else {
			this.backend.drawIndirect(commands, drawCount);
		}
	}

	public void close() {
		if (!this.isClosed) {
			this.isClosed = true;
			if (this.pushedDebugGroups > 0) {
				throw new IllegalStateException("Render pass had debug groups left open!");
			}

			this.onFinish.run();
		}
	}

	public record Draw<T>(
		int slot,
		GpuBuffer vertexBuffer,
		@Nullable GpuBuffer indexBuffer,
		@Nullable IndexType indexType,
		int firstIndex,
		int indexCount,
		int baseVertex,
		@Nullable BiConsumer<T, RenderPass.UniformUploader> uniformUploaderConsumer
	) {
		public Draw(
			final int slot,
			final GpuBuffer vertexBuffer,
			final GpuBuffer indexBuffer,
			final IndexType indexType,
			final int firstIndex,
			final int indexCount,
			final int baseVertex
		) {
			this(slot, vertexBuffer, indexBuffer, indexType, firstIndex, indexCount, baseVertex, null);
		}
	}

	public record RenderArea(int x, int y, int width, int height) {
		public boolean fillsTexture(final GpuTextureView texture) {
			return this.x == 0 && this.y == 0 && this.width == texture.getWidth(0) && this.height == texture.getHeight(0);
		}
	}

	public interface UniformUploader {
		void upload(String name, GpuBufferSlice buffer);
	}
}