package net.minecraft.client;

import com.mojang.blaze3d.buffers.GpuBuffer;
import com.mojang.blaze3d.pipeline.RenderTarget;
import com.mojang.blaze3d.platform.NativeImage;
import com.mojang.blaze3d.systems.CommandEncoder;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.blaze3d.textures.GpuTexture;
import com.mojang.logging.LogUtils;
import java.io.File;
import java.util.function.Consumer;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.ChatFormatting;
import net.minecraft.Util;
import net.minecraft.network.chat.ClickEvent;
import net.minecraft.network.chat.Component;
import net.minecraft.util.ARGB;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class Screenshot {
	private static final Logger LOGGER = LogUtils.getLogger();
	public static final String SCREENSHOT_DIR = "screenshots";

	/**
	 * Saves a screenshot in the game directory with a time-stamped filename.
	 */
	public static void grab(File gameDirectory, RenderTarget renderTarget, Consumer<Component> messageConsumer) {
		grab(gameDirectory, null, renderTarget, 1, messageConsumer);
	}

	public static void grab(File gameDirectory, @Nullable String filename, RenderTarget renderTarget, int downscaleFactor, Consumer<Component> messageConsumer) {
		takeScreenshot(
			renderTarget,
			downscaleFactor,
			nativeImage -> {
				File file2 = new File(gameDirectory, "screenshots");
				file2.mkdir();
				File file3;
				if (filename == null) {
					file3 = getFile(file2);
				} else {
					file3 = new File(file2, filename);
				}

				Util.ioPool()
					.execute(
						() -> {
							try {
								NativeImage exception = nativeImage;

								try {
									nativeImage.writeToFile(file3);
									Component component = Component.literal(file3.getName())
										.withStyle(ChatFormatting.UNDERLINE)
										.withStyle(style -> style.withClickEvent(new ClickEvent.OpenFile(file3.getAbsoluteFile())));
									messageConsumer.accept(Component.translatable("screenshot.success", component));
								} catch (Throwable var7) {
									if (nativeImage != null) {
										try {
											exception.close();
										} catch (Throwable var6) {
											var7.addSuppressed(var6);
										}
									}

									throw var7;
								}

								if (nativeImage != null) {
									nativeImage.close();
								}
							} catch (Exception var8) {
								LOGGER.warn("Couldn't save screenshot", (Throwable)var8);
								messageConsumer.accept(Component.translatable("screenshot.failure", var8.getMessage()));
							}
						}
					);
			}
		);
	}

	public static void takeScreenshot(RenderTarget renderTarget, Consumer<NativeImage> writer) {
		takeScreenshot(renderTarget, 1, writer);
	}

	public static void takeScreenshot(RenderTarget renderTarget, int downscaleFactor, Consumer<NativeImage> writer) {
		int i = renderTarget.width;
		int j = renderTarget.height;
		GpuTexture gpuTexture = renderTarget.getColorTexture();
		if (gpuTexture == null) {
			throw new IllegalStateException("Tried to capture screenshot of an incomplete framebuffer");
		} else if (i % downscaleFactor == 0 && j % downscaleFactor == 0) {
			GpuBuffer gpuBuffer = RenderSystem.getDevice().createBuffer(() -> "Screenshot buffer", 9, i * j * gpuTexture.getFormat().pixelSize());
			CommandEncoder commandEncoder = RenderSystem.getDevice().createCommandEncoder();
			RenderSystem.getDevice().createCommandEncoder().copyTextureToBuffer(gpuTexture, gpuBuffer, 0, () -> {
				try (GpuBuffer.MappedView mappedView = commandEncoder.mapBuffer(gpuBuffer, true, false)) {
					int l = j / downscaleFactor;
					int m = i / downscaleFactor;
					NativeImage nativeImage = new NativeImage(m, l, false);

					for (int n = 0; n < l; n++) {
						for (int o = 0; o < m; o++) {
							if (downscaleFactor == 1) {
								int p = mappedView.data().getInt((o + n * i) * gpuTexture.getFormat().pixelSize());
								nativeImage.setPixelABGR(o, j - n - 1, p | 0xFF000000);
							} else {
								int p = 0;
								int q = 0;
								int r = 0;

								for (int s = 0; s < downscaleFactor; s++) {
									for (int t = 0; t < downscaleFactor; t++) {
										int u = mappedView.data().getInt((o * downscaleFactor + s + (n * downscaleFactor + t) * i) * gpuTexture.getFormat().pixelSize());
										p += ARGB.red(u);
										q += ARGB.green(u);
										r += ARGB.blue(u);
									}
								}

								int s = downscaleFactor * downscaleFactor;
								nativeImage.setPixelABGR(o, l - n - 1, ARGB.color(255, p / s, q / s, r / s));
							}
						}
					}

					writer.accept(nativeImage);
				}

				gpuBuffer.close();
			}, 0);
		} else {
			throw new IllegalArgumentException("Image size is not divisible by downscale factor");
		}
	}

	/**
	 * Creates a unique PNG file in the given directory named by a timestamp. Handles cases where the timestamp alone is not enough to create a uniquely named file, though it still might suffer from an unlikely race condition where the filename was unique when this method was called, but another process or thread created a file at the same path immediately after this method returned.
	 */
	private static File getFile(File gameDirectory) {
		String string = Util.getFilenameFormattedDateTime();
		int i = 1;

		while (true) {
			File file = new File(gameDirectory, string + (i == 1 ? "" : "_" + i) + ".png");
			if (!file.exists()) {
				return file;
			}

			i++;
		}
	}
}