threedom/threedom-gl/src/main/kotlin/me/eater/threedom/gl/ShaderProgram.kt

package me.eater.threedom.gl

import me.eater.threedom.gl.GL.stack
import me.eater.threedom.gl.dom.ICamera
import me.eater.threedom.gl.exception.ShaderProgramLinkingException
import me.eater.threedom.gl.texture.Texture
import me.eater.threedom.utils.joml.toFloat
import org.joml.*
import org.lwjgl.opengl.GL20.*
import org.lwjgl.system.MemoryUtil
import kotlin.properties.ReadWriteProperty
import kotlin.reflect.KProperty

open class ShaderProgram(private val programId: Int = glCreateProgram()) {
    private val uniforms: MutableMap<String, Int> = mutableMapOf()
    val queuedUniformValues: MutableMap<Int, () -> Unit> = mutableMapOf()
    private val textures: MutableMap<Int, Int> = mutableMapOf()
    private var textureIndex: Int = 0

    var model by mat4()
    var normalModel by mat3()
    var view by mat4()
    var projection by mat4()

    private fun getUniformLocation(name: String) = uniforms.getOrPut(name) {
        glGetUniformLocation(programId, name)
    }

    protected fun sampler2D(
        index: Int = textureIndex++,
        name: String? = null
    ): ReadWriteProperty<ShaderProgram, Texture?> {
        var setUniform = false

        return object : ReadWriteProperty<ShaderProgram, Texture?> {
            override fun getValue(thisRef: ShaderProgram, property: KProperty<*>): Texture? {
                return Texture.getTexture(textures[index] ?: return null)
            }

            override fun setValue(thisRef: ShaderProgram, property: KProperty<*>, value: Texture?) {
                textures[index] = value?.id ?: 0

                if (setUniform) {
                    return
                }

                val loc = getUniformLocation(name ?: property.name)
                if (loc == -1) {
                    return
                }

                if (glGetInteger(GL_CURRENT_PROGRAM) == thisRef.programId) {
                    glUniform1i(loc, index)
                } else {
                    queuedUniformValues[loc] = {
                        glUniform1i(loc, index)
                    }
                }

                setUniform = true
            }
        }
    }

    protected fun <T> uniform(
        name: String? = null,
        setter: (location: Int, value: T) -> Unit,
        getter: (programId: Int, location: Int) -> T?
    ) = UniformProperty(name, setter, getter)

    open class UniformProperty<T>(
        val name: String? = null,
        val setter: (location: Int, value: T) -> Unit,
        val getter: (programId: Int, location: Int) -> T?
    ) {
        operator fun getValue(thisRef: ShaderProgram, property: KProperty<*>): T? {
            val loc = thisRef.getUniformLocation(name ?: property.name)
            if (loc == -1) {
                return null
            }

            return getter(thisRef.programId, loc)
        }

        operator fun setValue(thisRef: ShaderProgram, property: KProperty<*>, value: T?) {
            val valueNonNull = value ?: return

            val loc = thisRef.getUniformLocation(name ?: property.name)
            if (loc == -1) {
                return
            }


            if (glGetInteger(GL_CURRENT_PROGRAM) == thisRef.programId) {
                setter(loc, valueNonNull)
            } else {
                thisRef.queuedUniformValues[loc] = {
                    setter(loc, valueNonNull)
                }
            }
        }
    }

    protected fun float(name: String? = null): UniformProperty<Float> =
        uniform(name, ::glUniform1f, ::glGetUniformf)

    protected fun int(name: String? = null): UniformProperty<Int> =
        uniform(name, ::glUniform1i, ::glGetUniformi)

    protected fun vec2(name: String? = null): UniformProperty<Vector2fc> =
        uniform(name, { loc, value ->
            val floatArray = MemoryUtil.memAllocFloat(2)
            value.get(floatArray)
            glUniform2fv(loc, floatArray)
            MemoryUtil.memFree(floatArray)
        }, { programId, location ->
            val buffer = FloatArray(2)
            glGetUniformfv(programId, location, buffer)
            Vector2f().set(buffer)
        })

    protected fun vec3(name: String? = null): UniformProperty<Vector3fc> =
        uniform(name, { loc, value ->
            val floatArray = MemoryUtil.memAllocFloat(3)
            value.get(floatArray)
            glUniform3fv(loc, floatArray)
            MemoryUtil.memFree(floatArray)
        }, { programId, location ->
            val buffer = FloatArray(3)
            glGetUniformfv(programId, location, buffer)
            Vector3f().set(buffer)
        })

    protected fun vec4(name: String? = null): UniformProperty<Vector4fc> =
        uniform(name, { loc, value ->
            val floatArray = MemoryUtil.memAllocFloat(4)
            value.get(floatArray)
            glUniform4fv(loc, floatArray)
            MemoryUtil.memFree(floatArray)
        }, { programId, location ->
            val buffer = FloatArray(4)
            glGetUniformfv(programId, location, buffer)
            Vector4f().set(buffer)
        })

    protected fun mat3(name: String? = null): UniformProperty<Matrix3fc> =
        uniform(name, { loc, value ->
            val floatArray = FloatArray(3 * 3)
            value.get(floatArray)
            glUniformMatrix3fv(loc, false, floatArray)
        }, { programId, location ->
            val buffer = FloatArray(3 * 3)
            glGetUniformfv(programId, location, buffer)
            Matrix3f().set(buffer)
        })

    protected fun bool(name: String? = null): UniformProperty<Boolean> =
        uniform(name, { loc, value ->
            glUniform1i(loc, if (value) 1 else 0)
        }, { programId, location ->
            glGetUniformi(programId, location) == 1
        })

    protected fun mat4(name: String? = null): UniformProperty<Matrix4fc> =
        uniform(name, { loc, value ->
            val floatArray = FloatArray(4 * 4)
            value.get(floatArray)
            glUniformMatrix4fv(loc, false, floatArray)
        }, { programId, location ->
            val buffer = FloatArray(4 * 4)
            glGetUniformfv(programId, location, buffer)
            Matrix4f().set(buffer)
        })

    fun enable() {
        glUseProgram(programId)

        for ((_, queued) in queuedUniformValues) {
            queued()
        }

        for ((index, texture) in textures) {
            glActiveTexture(GL_TEXTURE0 + index)
            glBindTexture(GL_TEXTURE_2D, texture)
        }

        queuedUniformValues.clear()
    }

    fun delete() {
        glDeleteProgram(programId)
    }

    fun disable() {
        glUseProgram(0)
    }

    fun uses(name: String): Boolean {
        return getUniformLocation(name) != -1
    }



    fun camera(renderer: ICamera<*>) {
        this.projection = renderer.projection.toFloat()
        this.view = renderer.view.toFloat()
    }

    constructor(vertexShader: Shader, fragmentShader: Shader) : this() {
        glAttachShader(programId, vertexShader.shaderId)
        glAttachShader(programId, fragmentShader.shaderId)

        glLinkProgram(programId)
        val success = stack.mallocInt(1)
        glGetProgramiv(programId, GL_LINK_STATUS, success)
        if (success.get(0) == 0) {
            throw ShaderProgramLinkingException(glGetProgramInfoLog(programId))
        }
    }
}