package speiger.src.coreengine.platform.graphics.opengl.core;

import java.util.List;
import java.util.Objects;

import org.lwjgl.opengl.GL30;
import org.lwjgl.opengl.GL31;
import org.lwjgl.opengl.GL42;
import org.lwjgl.opengl.GL43;
import org.lwjgl.opengl.GL45;

import speiger.src.collections.objects.lists.ObjectArrayList;
import speiger.src.coreengine.platform.graphics.api.buffer.VertexBuffer;
import speiger.src.coreengine.platform.graphics.api.buffer.states.BufferType;
import speiger.src.coreengine.platform.graphics.api.compute.BarrierType;
import speiger.src.coreengine.platform.graphics.api.compute.ComputePipeline;
import speiger.src.coreengine.platform.graphics.api.core.ComputeCommandBuffer;
import speiger.src.coreengine.platform.graphics.api.sampler.Sampler;
import speiger.src.coreengine.platform.graphics.api.texture.Texture;
import speiger.src.coreengine.platform.graphics.api.utils.AccessType;
import speiger.src.coreengine.platform.graphics.opengl.buffer.GLVertexBuffer;
import speiger.src.coreengine.platform.graphics.opengl.sampler.GLSampler;
import speiger.src.coreengine.platform.graphics.opengl.shader.ShaderInstance;
import speiger.src.coreengine.platform.graphics.opengl.texture.GLTexture;
import speiger.src.coreengine.platform.graphics.opengl.utils.GLStates;
import speiger.src.coreengine.platform.graphics.opengl.utils.GLUtils;

public class GLRecordingComputeCommandBuffer implements ComputeCommandBuffer {
	List<Runnable> tasks = new ObjectArrayList<>();
	boolean recording = false;
	boolean removed = false;
	GLGraphicsDevice device;
	
	ComputePipeline pipeline;
	ShaderInstance shader;
	
	public GLRecordingComputeCommandBuffer(GLGraphicsDevice device) {
		this.device = device;
	}

	protected void ensureDrawing() {
		if(removed) throw new IllegalStateException("CommandBuffer is removed");		
		if(!recording) throw new IllegalStateException("CommandBuffer isn't recording");
	}
	
	@Override
	public boolean isRemoved() {
		return removed;
	}
	
	@Override
	public void remove() {
		removed = true;
		clear();
		recording = false;
		tasks.clear();
	}
	
	@Override
	public int commandCount() {
		return tasks.size();
	}
	
	@Override
	public ComputeCommandBuffer begin() {
		if(removed) throw new IllegalStateException("CommandBuffer is removed");
		if(recording) throw new IllegalStateException("CommandBuffer already recording");
		recording = true;
		tasks.clear();
		return this;
	}
	
	@Override
	public ComputeCommandBuffer pipeline(ComputePipeline pipeline) {
		Objects.requireNonNull(pipeline);
		ensureDrawing();
		if(this.pipeline == pipeline) return this;
		this.pipeline = pipeline;
		shader = device.getShader(pipeline);
		if(shader == null) {
			//TODO implement logging
			return this;
		}
		tasks.add(new SetPipeline(shader, device.states));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer texture(int binding, Texture texture, Sampler sampler) {
		Objects.requireNonNull(texture);
		Objects.requireNonNull(sampler);
		ensureDrawing();
		Objects.requireNonNull(shader, "No Pipeline found");
		Objects.requireNonNull(shader.samplers().byIndex(binding), "Texture["+binding+"] binding not found");
		tasks.add(new SetTexture(binding, (GLTexture)texture, (GLSampler)sampler, device.states));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer texture(int binding, Texture texture, AccessType access) {
		Objects.requireNonNull(texture);
		Objects.requireNonNull(access);
		ensureDrawing();
		Objects.requireNonNull(shader, "No Pipeline found");
		Objects.requireNonNull(shader.buffers().byIndex(binding), "Texture["+binding+"] binding not found");
		tasks.add(new SetImageStorage(binding, (GLTexture)texture, access));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer uniform(int binding, VertexBuffer buffer) {
		Objects.requireNonNull(buffer);
		if(buffer.type() != BufferType.UNIFORM_BUFFER) throw new IllegalStateException("Buffer isn't a Uniform Buffer");
		ensureDrawing();
		Objects.requireNonNull(shader, "No Pipeline found");
		Objects.requireNonNull(shader.uniforms().byIndex(binding), "Uniform for binding ["+binding+"] doesn't exist");
		tasks.add(new SetUniform(binding, (GLVertexBuffer)buffer, -1, -1, false));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer uniform(int binding, VertexBuffer buffer, long offset, long size) {
		Objects.requireNonNull(buffer);
		if(buffer.type() != BufferType.UNIFORM_BUFFER) throw new IllegalStateException("Buffer isn't a Uniform Buffer");
		ensureDrawing();
		Objects.requireNonNull(shader, "No Pipeline found");
		Objects.requireNonNull(shader.uniforms().byIndex(binding), "Uniform for binding ["+binding+"] doesn't exist");
		tasks.add(new SetUniform(binding, (GLVertexBuffer)buffer, offset, size, false));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer storage(int binding, VertexBuffer buffer) {
		Objects.requireNonNull(buffer);
		if(buffer.type() != BufferType.SHADER_STORAGE_BUFFER) throw new IllegalStateException("Buffer isn't a Shader Storage Buffer");
		ensureDrawing();
		Objects.requireNonNull(shader, "No Pipeline found");
		Objects.requireNonNull(shader.storages().byIndex(binding), "Shader Storage for binding ["+binding+"] doesn't exist");
		tasks.add(new SetUniform(binding, (GLVertexBuffer)buffer, -1, -1, true));
		return this;
	}

	@Override
	public ComputeCommandBuffer storage(int binding, VertexBuffer buffer, long offset, long size) {
		Objects.requireNonNull(buffer);
		if(buffer.type() != BufferType.SHADER_STORAGE_BUFFER) throw new IllegalStateException("Buffer isn't a Shader Storage Buffer");
		ensureDrawing();
		Objects.requireNonNull(shader, "No Pipeline found");
		Objects.requireNonNull(shader.storages().byIndex(binding), "Shader Storage for binding ["+binding+"] doesn't exist");
		tasks.add(new SetUniform(binding, (GLVertexBuffer)buffer, offset, size, true));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer dispatch(int numGroupsX, int numGroupsY, int numGroupsZ) {
		Objects.requireNonNull(shader);
		if(numGroupsX < 1) throw new IllegalArgumentException("X has to be 1 or more");
		if(numGroupsY < 1) throw new IllegalArgumentException("Y has to be 1 or more");
		if(numGroupsZ < 1) throw new IllegalArgumentException("Z has to be 1 or more");
		tasks.add(() -> GL45.glDispatchCompute(numGroupsX, numGroupsY, numGroupsZ));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer barrier(BarrierType... types) {
		int result = 0;
		for(BarrierType type : types) {
			result |= GLUtils.toGL(type);
		}
		tasks.add(new SetBarrier(result));
		return this;
	}
	
	@Override
	public ComputeCommandBuffer end() {
		if(!recording) throw new IllegalStateException("Already Stopped Recording");
		clear();
		return this;
	}
	
	public void execute() {
		tasks.forEach(Runnable::run);
	}
	
	private void clear() {
		recording = false;
		pipeline = null;
		shader = null;
	}
	
	public record SetPipeline(ShaderInstance instance, GLStates states) implements Runnable {
		@Override
		public void run() {
			states.shaders.bind(instance);
		}
	}
	
	public record SetTexture(int binding, GLTexture texture, GLSampler sampler, GLStates states) implements Runnable {
		@Override
		public void run() {
			states.textures.bind(binding, texture.id());
			states.samplers.bind(binding, sampler.id());
		}
	}
	
	public record SetImageStorage(int binding, GLTexture texture, AccessType type) implements Runnable {
		@Override
		public void run() {
			GL45.glBindImageTexture(0, binding, texture.id(), false, 0, GLUtils.toGL(type), GLUtils.toGLInternal(texture.settings().internal()));
		}
	}
	
	public record SetUniform(int binding, GLVertexBuffer buffer, long offset, long size, boolean shader) implements Runnable {
		@Override
		public void run() {
			if(size < 0 || offset < 0) GL30.glBindBufferBase(shader ? GL43.GL_SHADER_STORAGE_BUFFER : GL31.GL_UNIFORM_BUFFER, binding, buffer.id());
			else GL30.glBindBufferRange(shader ? GL43.GL_SHADER_STORAGE_BUFFER : GL31.GL_UNIFORM_BUFFER, binding, buffer.id(), offset, size);
		}
	}
	
	public record SetBarrier(int barrierBits) implements Runnable {
		@Override
		public void run() {
			GL42.glMemoryBarrier(barrierBits);
		}
	}
}
