package com.mojang.blaze3d;

import com.mojang.blaze3d.buffers.BufferType;
import com.mojang.blaze3d.buffers.BufferUsage;
import com.mojang.blaze3d.buffers.GpuBuffer;
import com.mojang.blaze3d.pipeline.RenderTarget;
import com.mojang.blaze3d.systems.CommandEncoder;
import com.mojang.blaze3d.systems.RenderPass;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.blaze3d.textures.GpuTexture;
import com.mojang.blaze3d.textures.TextureFormat;
import com.mojang.blaze3d.vertex.VertexFormat;
import com.mojang.jtracy.TracyClient;
import java.util.OptionalInt;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.RenderPipelines;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class TracyFrameCapture implements AutoCloseable {
	private static final int MAX_WIDTH = 320;
	private static final int MAX_HEIGHT = 180;
	private static final int BYTES_PER_PIXEL = 4;
	private int targetWidth;
	private int targetHeight;
	private int width;
	private int height;
	@Nullable
	private GpuTexture frameBuffer;
	@Nullable
	private GpuBuffer pixelbuffer;
	private int lastCaptureDelay;
	private boolean capturedThisFrame;
	private TracyFrameCapture.Status status = TracyFrameCapture.Status.WAITING_FOR_CAPTURE;

	private void resize(int width, int height) {
		float f = (float)width / height;
		if (width > 320) {
			width = 320;
			height = (int)(320.0F / f);
		}

		if (height > 180) {
			width = (int)(180.0F * f);
			height = 180;
		}

		width = width / 4 * 4;
		height = height / 4 * 4;
		if (this.width != width || this.height != height) {
			this.width = width;
			this.height = height;
			if (this.frameBuffer != null) {
				this.frameBuffer.close();
			}

			this.frameBuffer = RenderSystem.getDevice().createTexture("Tracy Frame Capture", TextureFormat.RGBA8, width, height, 1);
			if (this.pixelbuffer != null) {
				this.pixelbuffer.close();
			}

			this.pixelbuffer = RenderSystem.getDevice()
				.createBuffer(() -> "Tracy Frame Capture buffer", BufferType.PIXEL_PACK, BufferUsage.STREAM_READ, width * height * 4);
		}
	}

	public void capture(RenderTarget renderTarget) {
		if (this.status == TracyFrameCapture.Status.WAITING_FOR_CAPTURE
			&& !this.capturedThisFrame
			&& renderTarget.getColorTexture() != null
			&& this.pixelbuffer != null
			&& this.frameBuffer != null) {
			this.capturedThisFrame = true;
			if (renderTarget.width != this.targetWidth || renderTarget.height != this.targetHeight) {
				this.targetWidth = renderTarget.width;
				this.targetHeight = renderTarget.height;
				this.resize(this.targetWidth, this.targetHeight);
			}

			this.status = TracyFrameCapture.Status.WAITING_FOR_COPY;
			CommandEncoder commandEncoder = RenderSystem.getDevice().createCommandEncoder();
			RenderSystem.AutoStorageIndexBuffer autoStorageIndexBuffer = RenderSystem.getSequentialBuffer(VertexFormat.Mode.QUADS);
			GpuBuffer gpuBuffer = autoStorageIndexBuffer.getBuffer(6);

			try (RenderPass renderPass = RenderSystem.getDevice().createCommandEncoder().createRenderPass(this.frameBuffer, OptionalInt.empty())) {
				renderPass.setPipeline(RenderPipelines.TRACY_BLIT);
				renderPass.setVertexBuffer(0, RenderSystem.getQuadVertexBuffer());
				renderPass.setIndexBuffer(gpuBuffer, autoStorageIndexBuffer.type());
				renderPass.bindSampler("InSampler", renderTarget.getColorTexture());
				renderPass.drawIndexed(0, 6);
			}

			commandEncoder.copyTextureToBuffer(this.frameBuffer, this.pixelbuffer, 0, () -> this.status = TracyFrameCapture.Status.WAITING_FOR_UPLOAD, 0);
			this.lastCaptureDelay = 0;
		}
	}

	public void upload() {
		if (this.status == TracyFrameCapture.Status.WAITING_FOR_UPLOAD && this.pixelbuffer != null) {
			this.status = TracyFrameCapture.Status.WAITING_FOR_CAPTURE;

			try (GpuBuffer.ReadView readView = RenderSystem.getDevice().createCommandEncoder().readBuffer(this.pixelbuffer)) {
				TracyClient.frameImage(readView.data(), this.width, this.height, this.lastCaptureDelay, true);
			}
		}
	}

	public void endFrame() {
		this.lastCaptureDelay++;
		this.capturedThisFrame = false;
		TracyClient.markFrame();
	}

	public void close() {
		if (this.frameBuffer != null) {
			this.frameBuffer.close();
		}

		if (this.pixelbuffer != null) {
			this.pixelbuffer.close();
		}
	}

	@Environment(EnvType.CLIENT)
	static enum Status {
		WAITING_FOR_CAPTURE,
		WAITING_FOR_COPY,
		WAITING_FOR_UPLOAD;
	}
}