SimpleJavaEngine/src/main/java/speiger/src/coreengine/rendering/utils/ScreenshotHandler.java

package speiger.src.coreengine.rendering.utils;

import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.io.File;
import java.nio.ByteBuffer;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

import javax.imageio.ImageIO;

import org.lwjgl.opengl.GL11;
import org.lwjgl.system.MemoryUtil;

import speiger.src.coreengine.rendering.input.window.IWindowListener;
import speiger.src.coreengine.rendering.input.window.Window;
import speiger.src.coreengine.utils.helpers.FileUtils;
import speiger.src.coreengine.utils.io.GifWriter;

public class ScreenshotHandler implements IWindowListener
{
	public static final DateTimeFormatter FILE_FORMAT = DateTimeFormatter.ofPattern("yyyy-MM-dd_HH.mm.ss");
	public static final DateTimeFormatter FOLDER_FORMAT = DateTimeFormatter.ofPattern("yyyy-MM");
	static final ThreadPoolExecutor FACTORY = (ThreadPoolExecutor)Executors.newFixedThreadPool(1);

	File folder;
	int width;
	int height;
	ByteBuffer buffer;
	GifRecorder recorder;
	Supplier<BufferedImage> imageProvider = this::getData;
	
	public ScreenshotHandler(File helper)
	{
		FACTORY.setKeepAliveTime(1, TimeUnit.SECONDS);
		FACTORY.allowCoreThreadTimeOut(true);
		folder = new File(helper, "Screenshots");
	}
	
	@Override
	public void onWindowChanged(Window window)
	{
		width = window.getWidth();
		height = window.getHeight();
		if(buffer != null)
		{
			MemoryUtil.memFree(buffer);
			buffer = null;
		}
		buffer = MemoryUtil.memAlloc(width * height * 4);	
	}
	
	public void screenShot()
	{
		ZonedDateTime time = ZonedDateTime.now();
		File subFolder = new File(folder, time.format(FOLDER_FORMAT));
		FileUtils.ensureFolder(subFolder);
		try
		{
			ImageIO.write(getScreenShot(), "png", new File(subFolder, "Screenshot_"+time.format(FILE_FORMAT)+".png"));
		}
		catch(Exception e)
		{
			e.printStackTrace();
		}
	}
	
	public boolean isRecording()
	{
		return recorder != null && !recorder.isFinished();
	}
	
	public void toggleRecording()
	{
		if(isRecording())
		{
			stopRecording();
		}
		else
		{
			startRecording();
		}
	}
	
	//Undefined time recording
	public void startRecording()
	{
		record("Capture_"+ZonedDateTime.now().format(FILE_FORMAT)+".gif", 15, Integer.MAX_VALUE);
	}
	
	public void record(String fileName, int fps, int time)
	{
		if(recorder == null)
		{
			File subFolder = new File(folder, ZonedDateTime.now().format(FOLDER_FORMAT));
			FileUtils.ensureFolder(subFolder);
			recorder = new GifRecorder(new File(subFolder, fileName), fps, time);
		}
	}
	
	public void stopRecording()
	{
		if(recorder != null)
		{
			recorder.finishPic();
			recorder = null;
		}
	}
	
	public void update()
	{
		if(recorder != null)
		{
			recorder.tick(imageProvider);
			if(recorder.isFinished())
			{
				recorder.finishPic();
				recorder = null;
			}
		}
	}
	
	private BufferedImage getData()
	{
		buffer.clear();
		GL11.glReadBuffer(GL11.GL_FRONT);
		GL11.glReadPixels(0, 0, width, height, GL11.GL_RGBA, GL11.GL_UNSIGNED_BYTE, buffer);
		BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
		int[] data = ((DataBufferInt)image.getRaster().getDataBuffer()).getData();
		for(int x = 0;x < width;x++)
		{
			for(int y = 0;y < height;y++)
			{
				int i = (x + (width * y)) * 4;
				int r = buffer.get(i) & 0xFF;
				int g = buffer.get(i + 1) & 0xFF;
				int b = buffer.get(i + 2) & 0xFF;
				data[x + (width * (height - (y + 1)))] = (0xFF << 24) | (r << 16) | (g << 8) | b;
			}
		}
		return image;
	}
	
	public BufferedImage getScreenShot()
	{
		buffer.clear();
		GL11.glReadBuffer(GL11.GL_FRONT);
		GL11.glReadPixels(0, 0, width, height, GL11.GL_RGBA, GL11.GL_UNSIGNED_BYTE, buffer);
		BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
		int[] data = ((DataBufferInt)image.getRaster().getDataBuffer()).getData();
		for(int x = 0;x < width;x++)
		{
			for(int y = 0;y < height;y++)
			{
				int i = (x + (width * y)) * 4;
				int r = buffer.get(i) & 0xFF;
				int g = buffer.get(i + 1) & 0xFF;
				int b = buffer.get(i + 2) & 0xFF;
				data[x + (width * (height - (y + 1)))] = (0xFF << 24) | (r << 16) | (g << 8) | b;
			}
		}
		return image;
	}
	
	public void cleanup()
	{
		if(buffer != null)
		{
			MemoryUtil.memFree(buffer);
			buffer = null;
		}
	}
	
	public static class GifRecorder
	{
		GifWriter writer;
		long timeBetweenFrames;
		long lastTime = -1L;
		long freeTime;
		int totalPics;
		
		public GifRecorder(File file, int fps, int time)
		{
			int ms = 1000 / fps;
			writer = GifWriter.create(file, BufferedImage.TYPE_INT_ARGB, ms);
			timeBetweenFrames = ms * 1000000L;
			freeTime = ms * 1000000L;
			totalPics = fps * time;
		}
		
		public void tick(Supplier<BufferedImage> supply)
		{
			if(lastTime == -1L)
			{
				totalPics--;
				insert(supply.get());
				lastTime = System.nanoTime();
				return;
			}
			long time = System.nanoTime();
			long happend = time - lastTime;
			lastTime = time;
			freeTime -= happend;
			if(freeTime <= 0)
			{
				freeTime += timeBetweenFrames;
				insert(supply.get());
				totalPics--;
			}
		}
		
		private void insert(final BufferedImage newBuffer)
		{
			FACTORY.execute(new Runnable(){
				BufferedImage buffer = newBuffer;
				@Override
				public void run()
				{
					if(writer != null)
					{
						writer.insertPicture(buffer);
					}
				}
			});
		}
		
		public boolean isFinished()
		{
			return totalPics <= 0;
		}
		
		public void finishPic()
		{
			FACTORY.execute(new Runnable(){
				@Override
				public void run()
				{
					writer.close();
					writer = null;
				}
			});
		}
	}
}