package ru.dbotthepony.kstarbound.client.gl.shader

import com.github.benmanes.caffeine.cache.Cache
import com.github.benmanes.caffeine.cache.Caffeine
import com.github.benmanes.caffeine.cache.Scheduler
import it.unimi.dsi.fastutil.objects.Object2ObjectArrayMap
import it.unimi.dsi.fastutil.objects.ObjectArraySet
import it.unimi.dsi.fastutil.objects.Reference2ObjectFunction
import org.lwjgl.opengl.GL20.GL_FRAGMENT_SHADER
import org.lwjgl.opengl.GL20.GL_VERTEX_SHADER
import ru.dbotthepony.kommons.util.IStruct4f
import ru.dbotthepony.kommons.math.RGBAColor
import ru.dbotthepony.kommons.matrix.Matrix3f
import ru.dbotthepony.kstarbound.client.StarboundClient
import ru.dbotthepony.kstarbound.client.gl.GLTexture2D
import ru.dbotthepony.kstarbound.client.gl.vertex.StreamVertexBuilder
import ru.dbotthepony.kstarbound.client.gl.vertex.VertexAttributeType
import ru.dbotthepony.kstarbound.client.gl.vertex.VertexAttributes
import ru.dbotthepony.kstarbound.client.render.RenderConfig
import java.time.Duration
import java.util.Collections
import java.util.EnumSet
import java.util.function.Function

class UberShader private constructor(
	fragment: GLShader,
	vertex: GLShader,
	attributes: VertexAttributes,
	val flags: Set<Flag>,
) : GLShaderProgram(listOf(fragment, vertex), attributes) {
	enum class Flag {
		LIGHT_MAP,
		NEEDS_SCREEN_SIZE;
	}

	var viewMatrix: Matrix3f by F3x3Uniform("viewMatrix")
	var modelMatrix: Matrix3f by F3x3Uniform("modelMatrix")
	var colorMultiplier: IStruct4f by F4Uniform("colorMultiplier")

	var texture0 by IUniform("texture0", VertexAttributeType.UV !in attributes)

	var lightmapTexture by IUniform("lightmap", Flag.LIGHT_MAP !in flags)
	var lightmapUV by F4Uniform("lightmapUV", Flag.LIGHT_MAP !in flags)
	var screenSize by F2Uniform("screenSize", Flag.NEEDS_SCREEN_SIZE !in flags)

	val builder = StreamVertexBuilder(attributes)

	init {
		viewMatrix = Matrix3f.identity()
		modelMatrix = Matrix3f.identity()
		colorMultiplier = RGBAColor.WHITE
	}

	private val textureConfigs: Cache<GLTexture2D, RenderConfig<UberShader>> = Caffeine.newBuilder()
		.weakValues()
		.weakKeys()
		.expireAfterAccess(Duration.ofMinutes(1))
		.build()

	fun config(texture: GLTexture2D): RenderConfig<UberShader> {
		return textureConfigs.get(texture) {
			object : RenderConfig<UberShader>(this) {
				override fun setup() {
					super.setup()

					client.textures2D[0] = it
				}
			}
		}
	}

	class Builder {
		private val directives = Object2ObjectArrayMap<String, String>()
		private val attributes = ObjectArraySet<VertexAttributeType>()
		private val flags = EnumSet.noneOf(Flag::class.java)

		private fun addFlag(flag: Flag): Builder {
			if (flags.add(flag)) {
				directives[flag.name] = ""
			}

			return this
		}

		init {
			attributes.add(VertexAttributeType.POSITION)
		}

		private var nextAttributeIndex = 1

		private fun attribute(type: VertexAttributeType, name: String = type.name): Builder {
			if (attributes.add(type)) {
				directives[name] = nextAttributeIndex.toString()
				nextAttributeIndex += type.type.width
			}

			return this
		}

		fun withTexture(): Builder {
			return attribute(VertexAttributeType.UV, "TEXTURE")
		}

		fun withColor(): Builder {
			return attribute(VertexAttributeType.COLOR)
		}

		fun withHueShift(): Builder {
			return attribute(VertexAttributeType.HUE_SHIFT)
		}

		fun withLightMap(): Builder {
			addFlag(Flag.NEEDS_SCREEN_SIZE)
			return addFlag(Flag.LIGHT_MAP)
		}

		fun build(): UberShader {
			val fragment = GLShader("#version 450\n" +
				                        directives.entries.joinToString("\n") { "#define ${it.key}" + (if (it.value != "") " " + it.value else "") } +
				                        fragment, GL_FRAGMENT_SHADER)

			val vertex = GLShader("#version 450\n" +
				                        directives.entries.joinToString("\n") { "#define ${it.key}" + (if (it.value != "") " " + it.value else "") } +
				                        vertex, GL_VERTEX_SHADER)

			val attributes = VertexAttributes.of(attributes)

			return UberShader(fragment, vertex, attributes, Collections.unmodifiableSet(EnumSet.copyOf(flags)))
		}
	}

	companion object {
		private val fragment by lazy { StarboundClient.readInternal("shaders/fragment/configurable.fsh") }
		private val vertex by lazy { StarboundClient.readInternal("shaders/vertex/configurable.vsh") }
	}
}
