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

import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.function.IntPredicate;
import java.util.function.Supplier;

import org.lwjgl.opengl.GL;
import org.lwjgl.opengl.GL11;
import org.lwjgl.opengl.GL12;
import org.lwjgl.opengl.GL20;
import org.lwjgl.opengl.GL31;
import org.lwjgl.opengl.GL33;
import org.lwjgl.opengl.GL41;
import org.lwjgl.opengl.GL43;
import org.lwjgl.opengl.GL45;
import org.lwjgl.opengl.GL46;
import org.lwjgl.opengl.GLCapabilities;
import org.lwjgl.system.MemoryUtil;

import speiger.src.collections.ints.lists.IntArrayList;
import speiger.src.collections.ints.lists.IntList;
import speiger.src.collections.ints.maps.impl.hash.Int2ObjectLinkedOpenHashMap;
import speiger.src.collections.ints.maps.impl.hash.Int2ObjectOpenHashMap;
import speiger.src.collections.ints.maps.interfaces.Int2ObjectMap;
import speiger.src.collections.objects.maps.impl.hash.Object2IntOpenHashMap;
import speiger.src.collections.objects.maps.impl.hash.Object2ObjectOpenHashMap;
import speiger.src.collections.objects.maps.interfaces.Object2IntMap;
import speiger.src.collections.objects.maps.interfaces.Object2ObjectMap;
import speiger.src.collections.objects.utils.ObjectIterables;
import speiger.src.coreengine.assets.api.IAsset;
import speiger.src.coreengine.assets.api.IAssetProvider;
import speiger.src.coreengine.assets.api.ID;
import speiger.src.coreengine.assets.api.MultiAsset;
import speiger.src.coreengine.assets.manager.AssetManager;
import speiger.src.coreengine.platform.graphics.api.buffer.VertexBuffer;
import speiger.src.coreengine.platform.graphics.api.buffer.states.BufferState;
import speiger.src.coreengine.platform.graphics.api.buffer.states.BufferType;
import speiger.src.coreengine.platform.graphics.api.compute.ComputePipeline;
import speiger.src.coreengine.platform.graphics.api.core.GraphicsCommandBuffer;
import speiger.src.coreengine.platform.graphics.api.core.GraphicsDevice;
import speiger.src.coreengine.platform.graphics.api.core.TimeQueryPool;
import speiger.src.coreengine.platform.graphics.api.mesh.DynamicMesh;
import speiger.src.coreengine.platform.graphics.api.mesh.Mesh;
import speiger.src.coreengine.platform.graphics.api.mesh.Mesh.LayoutInfo;
import speiger.src.coreengine.platform.graphics.api.sampler.SamplerSettings;
import speiger.src.coreengine.platform.graphics.api.shader.ShaderPipeline;
import speiger.src.coreengine.platform.graphics.api.shader.VertexBinding;
import speiger.src.coreengine.platform.graphics.api.shader.states.Bindings.BufferBinding;
import speiger.src.coreengine.platform.graphics.api.shader.states.Bindings.BufferMode;
import speiger.src.coreengine.platform.graphics.api.shader.states.Bindings.TextureBinding;
import speiger.src.coreengine.platform.graphics.api.shader.states.Bindings.TextureMode;
import speiger.src.coreengine.platform.graphics.api.shader.states.ShaderType;
import speiger.src.coreengine.platform.graphics.api.texture.TextureSettings;
import speiger.src.coreengine.platform.graphics.api.texture.states.SwizzleMask;
import speiger.src.coreengine.platform.graphics.api.texture.states.TextureType;
import speiger.src.coreengine.platform.graphics.api.utils.ExecutionType;
import speiger.src.coreengine.platform.graphics.api.vertex.VertexLayout.Element;
import speiger.src.coreengine.platform.graphics.opengl.buffer.GLVertexBuffer;
import speiger.src.coreengine.platform.graphics.opengl.mesh.GLDynamicMesh;
import speiger.src.coreengine.platform.graphics.opengl.mesh.GLMesh;
import speiger.src.coreengine.platform.graphics.opengl.mesh.VertexArrayCache;
import speiger.src.coreengine.platform.graphics.opengl.sampler.GLSampler;
import speiger.src.coreengine.platform.graphics.opengl.shader.ShaderCache;
import speiger.src.coreengine.platform.graphics.opengl.shader.ShaderInstance;
import speiger.src.coreengine.platform.graphics.opengl.shader.ShaderInstance.BufferObject;
import speiger.src.coreengine.platform.graphics.opengl.shader.ShaderInstance.ShaderStorage;
import speiger.src.coreengine.platform.graphics.opengl.shader.ShaderInstance.TextureObject;
import speiger.src.coreengine.platform.graphics.opengl.texture.GLTexture;
import speiger.src.coreengine.platform.graphics.opengl.utils.GLFunctions;
import speiger.src.coreengine.platform.graphics.opengl.utils.GLStates;
import speiger.src.coreengine.platform.graphics.opengl.utils.GLUtils;
import speiger.src.coreengine.platform.input.window.Window;

public class GLGraphicsDevice implements GraphicsDevice {
	static final Map<ID, String> IMPORTS = Object2ObjectMap.builder().<ID, String>map().synchronize();
	Object2IntMap<ID> shaderCache = new Object2IntOpenHashMap<>();
	Object2ObjectMap<ID, ShaderInstance> programCache = new Object2ObjectOpenHashMap<>();
	VertexArrayCache vertexCache = new VertexArrayCache();
	AssetManager manager;
	GLCapabilities capabilities;
	GLStates states = new GLStates();
	ShaderCache cache = new ShaderCache();
	GLCommandQueue queue;
	Window owner;
	
	public GLGraphicsDevice(Window owner, AssetManager manager, GLCapabilities capbilities) {
		this.owner = owner;
		this.manager = manager;
		this.capabilities = capbilities;
		GL.setCapabilities(capbilities);
		this.queue = new GLCommandQueue(this);
	}
	
	@Override
	public Window window() {
		return owner;
	}
	
	@Override
	public GLSurface createSurface() {
		return new GLSurface(owner, this);
	}
	
	@Override
	public GraphicsCommandBuffer createCommandBuffer(ExecutionType type) {
		return type == ExecutionType.IMMEDIATE ? new GLImmidateCommandBuffer(this) : new GLRecordingCommandBuffer(this);
	}

	@Override
	public GLCommandQueue queue() {
		return queue;
	}
	
	@Override
	public GLVertexBuffer createBuffer(BufferType type, BufferState state) {
		return new GLVertexBuffer(Objects.requireNonNull(type), Objects.requireNonNull(state));
	}
	
	@Override
	public GLVertexBuffer createBuffer(BufferType type, BufferState state, int allocatedBytes) {
		if(allocatedBytes <= 0) throw new IllegalArgumentException("Allocated bytes are invalid");
		return new GLVertexBuffer(Objects.requireNonNull(type), Objects.requireNonNull(state), allocatedBytes);
	}
	
	@Override
	public GLTexture createTexture(TextureSettings settings, int width, int height) {
		if(width < 0) throw new IllegalArgumentException("Width ["+width+"] is invalid");
		if(height < 0) throw new IllegalArgumentException("Height ["+height+"] is invalid");
		if(settings.type() != TextureType.TEXTURE_2D) throw new IllegalStateException("Currently only 2D textures are supported");
		int id = GLFunctions.createTexture(settings.type());
		settings.stencilType().ifPresent(T -> GL45.glTextureParameteri(id, GL43.GL_DEPTH_STENCIL_TEXTURE_MODE, GLUtils.toGL(T)));
		settings.baseLevel().ifPresent(T -> GL45.glTextureParameteri(id, GL12.GL_TEXTURE_BASE_LEVEL, T));
		settings.maxLevel().ifPresent(T -> GL45.glTextureParameteri(id, GL12.GL_TEXTURE_MAX_LEVEL, T));
		if(settings.hasSwizzle()) {
			List<Optional<SwizzleMask>> masks = settings.swizzle();
			GL45.glTextureParameterIiv(id, GL33.GL_TEXTURE_SWIZZLE_RGBA, new int[] {
					GLUtils.toGL(masks.get(0).orElse(SwizzleMask.RED)),
					GLUtils.toGL(masks.get(1).orElse(SwizzleMask.GREEN)),
					GLUtils.toGL(masks.get(2).orElse(SwizzleMask.BLUE)),
					GLUtils.toGL(masks.get(3).orElse(SwizzleMask.ALPHA)),
			});
		}
		GL45.glTextureStorage2D(id, 1, GLUtils.toGLInternal(settings.internal()), width, height);
		return new GLTexture(settings, id, width, height);
	}
	
	@Override
	public GLSampler createSampler(SamplerSettings settings) {
		return new GLSampler(Objects.requireNonNull(settings));
	}
	
	@Override
	public GLMesh createMesh(Mesh.Builder builder) {
		builder.validate();
		if(builder.isAutogenerated()) {
			Int2ObjectMap<LayoutInfo> info = builder.info();
			OptionalInt preAllocation = builder.preAllocated();
			Int2ObjectMap<VertexBuffer> buffers = new Int2ObjectLinkedOpenHashMap<>();
			for(Int2ObjectMap.Entry<LayoutInfo> entry : info.int2ObjectEntrySet()) {
				int key = entry.getIntKey();
				buffers.put(key, preAllocation.isPresent() ? createBuffer(BufferType.ARRAY_BUFFER, builder.getState(key), preAllocation.getAsInt()) : createBuffer(BufferType.ARRAY_BUFFER, builder.getState(key)));
			}
			return new GLMesh(info, buffers, createBuffer(BufferType.ELEMENT_BUFFER, builder.getState(-1)), builder.indecies(), true);
		}
		return new GLMesh(builder.info(), new Int2ObjectOpenHashMap<>(), null, null, false);
	}
	
	@Override
	public DynamicMesh createDynamicMesh(DynamicMesh.Builder builder) {
		if(builder.isAutogenerated()) {
			OptionalInt preAllocation = builder.preAllocated();
			Int2ObjectMap<VertexBuffer> buffers = new Int2ObjectLinkedOpenHashMap<>();
			for(Int2ObjectMap.Entry<BufferState> entry : builder.buffers().int2ObjectEntrySet()) {
				int key = entry.getIntKey();
				if(key == -1) continue;
				buffers.put(key, preAllocation.isPresent() ? createBuffer(BufferType.ARRAY_BUFFER, entry.getValue(), preAllocation.getAsInt()) : createBuffer(BufferType.ARRAY_BUFFER, entry.getValue()));
			}
			return new GLDynamicMesh(buffers, createBuffer(BufferType.ELEMENT_BUFFER, builder.getIndeciesBuffer()), builder.indecies(), true);

		}
		return new GLDynamicMesh(new Int2ObjectOpenHashMap<>(), null, null, false);
	}
	
	@Override
	public TimeQueryPool createQueryPool(int size) {
		return new GLTimeQueryPool(size);
	}
	
	@Override
	public void preloadPipelines(List<ShaderPipeline> graphics, List<ComputePipeline> compute) {
		try(IAssetProvider provider = manager.get()) {
			for(ShaderPipeline pipeline : graphics) {
				programCache.put(pipeline.id(), computeShaderProgram(pipeline, provider));
			}
			for(ComputePipeline pipeline : compute) {
				programCache.put(pipeline.id(), computeShaderProgram(pipeline, provider));
			}
		}
		catch(Exception e) {
			e.printStackTrace();
		}
	}
	
	public ShaderInstance getShader(ShaderPipeline pipeline) {
		return programCache.supplyIfAbsent(pipeline.id(), () -> computeShaderProgram(pipeline));
	}
	
	private ShaderInstance computeShaderProgram(ShaderPipeline pipeline) {
		try(IAssetProvider provider = manager.get()) {
			return computeShaderProgram(pipeline, provider);			
		}
		catch(Exception e) {
			e.printStackTrace();
			return null;
		}
	}
	
	public ShaderInstance getShader(ComputePipeline pipeline) {
		return programCache.supplyIfAbsent(pipeline.id(), () -> computeShaderProgram(pipeline));
	}
	
	private ShaderInstance computeShaderProgram(ComputePipeline pipeline) {
		try(IAssetProvider provider = manager.get()) {
			return computeShaderProgram(pipeline, provider);			
		}
		catch(Exception e) {
			e.printStackTrace();
			return null;
		}
	}
	
	
	private ShaderInstance computeShaderProgram(ComputePipeline pipeline, IAssetProvider provider) {
		MultiAsset assets = MultiAsset.combineNonNull(provider, pipeline.shader());
		if(assets == null) {
			//TODO log stuff
			return null;
		}
		
		int program = loadOrCreateProgram(pipeline.id(), assets::crc, T -> {
			int id = shaderCache.computeIfAbsent(pipeline.shader(), E -> computeShader(E, assets.get(0), ShaderType.COMPUTE, provider));
			if(id == -1) {
				System.out.println("Couldn't create a compute shader Id["+pipeline.shader()+"]");
				//TODO implement logging.
				//Failed to compile shader
				return false;
			}
			GL20.glAttachShader(T, id);
			GL20.glLinkProgram(T);
			GL20.glDetachShader(T, id);
			return true;
		});
		if(program == 0) return null;
		return generateShaderInstance(program, pipeline::textures, pipeline::uniforms);
	}
	
	private ShaderInstance computeShaderProgram(ShaderPipeline pipeline, IAssetProvider provider) {
		MultiAsset assets = MultiAsset.combineNonNull(provider, pipeline.shaders().values().toArray(ID[]::new));
		if(assets == null) {
			//TODO log stuff
			return null;
		}
		int program = loadOrCreateProgram(pipeline.id(), assets::crc, T -> {
			Map<ID, IAsset> map = Object2ObjectMap.builder().map();
			assets.forEach(map, (K, V) -> K.put(V.location(), V));
			IntList shaders = new IntArrayList();
			for(Entry<ShaderType, ID> entry : pipeline.shaders().entrySet()) {
				ShaderType type = entry.getKey();
				IAsset asset = map.get(entry.getValue());
				int id = shaderCache.computeIfAbsent(entry.getValue(), E -> computeShader(E, asset, type, provider));
				if(id == -1) {
					System.out.println("Couldn't create a shader. Type["+type+"], of Id["+asset.location()+"]");
					//TODO implement logging.
					//Failed to compile shader
					return false;
				}
				GL20.glAttachShader(T, id);
				shaders.add(id);
			}
			for(Element element : ObjectIterables.flatMap(pipeline.attributes(), VertexBinding::elements)) {
				GL20.glBindAttribLocation(T, element.index(), element.name());
			}
			GL20.glLinkProgram(T);
			for(int i = 0,m=shaders.size();i<m;i++) {
				GL20.glDetachShader(T, shaders.getInt(i));
			}
			return true;
		});
		if(program == 0) return null;
		return generateShaderInstance(program, pipeline::textures, pipeline::uniforms);
	}
	
	private ShaderInstance generateShaderInstance(int program, Supplier<List<TextureBinding>> textureProvider, Supplier<List<BufferBinding>> uniformProvider) {
		ShaderStorage<BufferObject> uniforms = new ShaderStorage<>();
		ShaderStorage<BufferObject> storages = new ShaderStorage<>();
		ShaderStorage<TextureObject> samplers = new ShaderStorage<>();
		ShaderStorage<TextureObject> buffers = new ShaderStorage<>();
		for(TextureBinding binding : textureProvider.get()) {
			int location = GL20.glGetUniformLocation(program, binding.name());
			if(location == -1) {
				System.out.println("Couldn't find Location ["+binding.name()+"] in Program ["+program+"]");
				//TODO Warn that location wasn't found
				continue;
			}
			(binding.mode() == TextureMode.SAMPLER ? samplers : buffers).add(binding.name(), binding.slot(), new TextureObject(binding.slot()));
		}
		for(BufferBinding binding : uniformProvider.get()) {
			int location = GL31.glGetUniformBlockIndex(program, binding.name());
			if(location == GL31.GL_INVALID_INDEX) {
				System.out.println("Couldn't find Location ["+binding.name()+"] in Program ["+program+"]");
				//TODO Warn that location wasn't found
				continue;
			}
			GL31.glUniformBlockBinding(program, location, binding.slot());
			(binding.mode() == BufferMode.UNIFORM ? uniforms : storages).add(binding.name(), binding.slot(), new BufferObject(binding.slot()));
		}
		return new ShaderInstance(program, uniforms, storages, samplers, buffers);
	}
	
	private int computeShader(ID id, IAsset asset, ShaderType type, IAssetProvider provider) {
		StringBuilder builder = new StringBuilder(2048);
		try {
			for(String line : asset.lines()) {
				if(line.startsWith("//")) continue;
				if(line.startsWith("#import<") && line.endsWith(">")) builder.append(importLines(provider, ID.of(line.substring(8, line.length()-1))));
				else builder.append(line);
				builder.append("\n");
			};
		}
		catch(Exception e) {
			e.printStackTrace();
			return -1;
		}
		int shader = GL20.glCreateShader(GLUtils.toGL(type));
		GL20.glShaderSource(shader, builder);
		GL20.glCompileShader(shader);
		if(GL20.glGetShaderi(shader, GL20.GL_COMPILE_STATUS) == 0) {
			//TODO improve logging
			System.out.println(GL20.glGetShaderInfoLog(shader, GL20.glGetShaderi(shader, GL20.GL_INFO_LOG_LENGTH)));
			GL20.glDeleteShader(shader);
			return -1;
		}
		return shader;
	}
	
	private int loadOrCreateProgram(ID pipeline, Supplier<String> crc, IntPredicate callback) {
		int id = GL20.glCreateProgram();
		boolean fail = true;
		if((fail = !loadFromCache(id, pipeline, crc) && callback.test(id))) {
			if(GL20.glGetProgrami(id, GL20.GL_LINK_STATUS) == GL11.GL_TRUE) {
				cache.store(pipeline, crc.get(), getProgramBytes(id));
				fail = false;
			}
		}
		if(fail) {
			GL20.glDeleteProgram(id);
			return 0;
		}
		String error = GL20.glGetProgramInfoLog(id);
		if(!error.isBlank()) {
			System.out.println("Error: "+error);
		}
		return id;
	}
	
	private byte[] getProgramBytes(int programId) {
		ByteBuffer buffer = MemoryUtil.memAlloc(GL20.glGetProgrami(programId, GL41.GL_PROGRAM_BINARY_LENGTH) + 4);
		int[] format = new int[1];
		GL41.glGetProgramBinary(programId, null, format, buffer.position(4));
		buffer.putInt(0, format[0]).flip();
		byte[] data = new byte[buffer.remaining()];
		buffer.get(data);
		MemoryUtil.memFree(buffer);
		return data;
	}
	
	private boolean loadFromCache(int programId, ID pipeline, Supplier<String> crc) {
		ByteBuffer data = cache.get(pipeline, crc);
		if(data == null) return false;
		GL41.glProgramBinary(programId, data.getInt(), data);
		MemoryUtil.memFree(data);
		return GL46.glGetProgrami(programId, GL46.GL_LINK_STATUS) == GL46.GL_TRUE;
	}
	
	public static String importLines(IAssetProvider provider, ID location) {
		return IMPORTS.computeIfAbsent(location, T -> {
			try { return ObjectIterables.filter(provider.get(T).lines(), E -> !E.startsWith("//")).reduce(new StringBuilder(), (K, V) -> K.append(V).append("\n")).toString(); }
			catch(Exception e) {
				e.printStackTrace();
				return "";
			}
		});
	}
}
