package speiger.src.coreengine.rendering.shader;

import java.nio.ByteBuffer;
import java.nio.file.attribute.FileTime;
import java.util.Map;
import java.util.function.IntPredicate;

import org.lwjgl.opengl.GL20;
import org.lwjgl.opengl.GL41;
import org.lwjgl.opengl.GL45;
import org.lwjgl.opengl.GL46;
import org.lwjgl.system.MemoryUtil;

import speiger.src.collections.objects.maps.interfaces.Object2ObjectMap;
import speiger.src.coreengine.assets.AssetLocation;
import speiger.src.coreengine.assets.base.IAsset;
import speiger.src.coreengine.assets.base.IAssetProvider;
import speiger.src.coreengine.math.ArrayUtil;
import speiger.src.coreengine.rendering.shader.uniform.UniformManager;
import speiger.src.coreengine.rendering.utils.GLStateTracker;
import speiger.src.coreengine.rendering.utils.values.IGLValue.IShaderType;
import speiger.src.coreengine.utils.io.GameLog;

public abstract class ShaderProgram {
	static final Map<AssetLocation, String> IMPORTS = Object2ObjectMap.builder().<AssetLocation, String>map().synchronize();
	protected final UniformManager uniforms = new UniformManager(this);
	private int id;
	
	public boolean isValid() { return id != 0; }
	public boolean isActive() { return GLStateTracker.instance().shaders.isShaderActive(this); }
	public int id() { return id; }
	public UniformManager getUniforms() { return uniforms; }
	
	protected void setId(int id) {
		this.id = id;
	}
	
	public void bind() {
		GLStateTracker.instance().shaders.bind(this);
		uniforms.bind();
	}
	
	public void remove() {
		if(id == 0) return;
		uniforms.remove();
		GL20.glDeleteProgram(id);
		id = 0;
	}
	
	public void validateProgram() {
		if(id == 0) return;
		uniforms.validate();
		GL46.glValidateProgram(id);
	}
	
	public static int loadOrGenerateCache(String location, FileTime newestResource, IntPredicate callback) {
		int id = GL20.glCreateProgram();
		boolean fail = true;
		if((fail = !loadFromBinary(id, location, newestResource)) && callback.test(id)) {
			if(GL46.glGetProgrami(id, GL46.GL_LINK_STATUS) == GL46.GL_TRUE) {
				if(ShaderCache.INSTANCE.hasCache()) ShaderCache.INSTANCE.storeInCache(location, getProgramBytes(id));
				fail = false;
			}
		}
		if(fail) {
			GL20.glDeleteProgram(id);
			return 0;
		}
		String error = GL45.glGetProgramInfoLog(id);
		if(!error.isBlank()) {
			System.out.println("Error: "+error);
		}
		return id;
	}
	
	private static boolean loadFromBinary(int programId, String location, FileTime newestResource) {
		ByteBuffer buffer = ShaderCache.INSTANCE.readFromCache(location, newestResource);
		if(buffer != null) {
			GL41.glProgramBinary(programId, buffer.getInt(), buffer);
			MemoryUtil.memFree(buffer);
			if(GL46.glGetProgrami(programId, GL46.GL_LINK_STATUS) == GL46.GL_TRUE) return true;
		}
		return false;
	}
	
	private static 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));
		return ArrayUtil.toArray(buffer.putInt(0, format[0]).flip(), MemoryUtil::memFree);
	}
	
	public static int loadShader(IAssetProvider provider, AssetLocation location, IShaderType type) {
		StringBuilder builder = new StringBuilder(2048);
		try(IAsset asset = provider.getAsset(location)) {
			for(String line : asset.lines()) {
				if(line.startsWith("//")) continue;
				if(line.startsWith("#import<") && line.endsWith(">")) {
					builder.append(importLines(provider, AssetLocation.of(line.substring(8, line.length()-1)))).append("\n");
					continue;
				}
				builder.append(line).append("\n");
			}
		}
		catch(Exception e) {
			e.printStackTrace();
			return -1;
		}
		int shaderID = GL20.glCreateShader(type.glValue());
		GL20.glShaderSource(shaderID, builder);
		GL20.glCompileShader(shaderID);
		if(GL20.glGetShaderi(shaderID, GL20.GL_COMPILE_STATUS) == 0) {
			GameLog.warn("Could not compile shader " + location);
			GameLog.error(GL20.glGetShaderInfoLog(shaderID, GL20.glGetShaderi(shaderID, GL20.GL_INFO_LOG_LENGTH)));
			GL20.glDeleteShader(shaderID);
			return -1;
		}
		return shaderID;
	}
	
	public static String importLines(IAssetProvider provider, AssetLocation location) {
		return IMPORTS.computeIfAbsent(location, T -> {
			try(IAsset asset = provider.getAsset(T)) {
				StringBuilder builder = new StringBuilder();
				for(String line : asset.lines()) {
					if(line.startsWith("//")) continue;
					builder.append(line).append("\n");
				}
				return builder.toString();
			}
			catch(Exception e) {
				e.printStackTrace();
				return "";
			}
		});
	}
}
