GlslBase.kt

package baaahs.glsl

import baaahs.*
import baaahs.geom.Vector2F
import baaahs.shows.GlslShow
import org.joml.Matrix4f
import org.lwjgl.BufferUtils
import org.lwjgl.glfw.GLFW.*
import org.lwjgl.glfw.GLFWErrorCallback
import org.lwjgl.glfw.GLFWFramebufferSizeCallback
import org.lwjgl.glfw.GLFWKeyCallback
import org.lwjgl.opengl.*
import org.lwjgl.opengl.GL31.*
import org.lwjgl.system.MemoryUtil.NULL
import org.lwjgl.system.MemoryUtil.memAddress
import java.nio.ByteBuffer
import kotlin.math.max
import kotlin.math.min
import kotlin.math.round
import kotlin.random.Random

fun main() {
    // On a Mac, it's necessary to start the JVM with this arg: `-XstartOnFirstThread`
    val renderer = GlslBase.manager.getRenderer(GlslShow.program) as JvmGlslRenderer

    renderer.addSurface(
        IdentifiedSurface(
            SheepModel.Panel("Panel"),
            600,
            List(600) { Vector2F(Random.nextFloat(), Random.nextFloat()) })
    )

    renderer.runStandalone()
}

interface ViewSettings {
    companion object {

        /**
         * The distance between the viewer's eyes and the screen in some distance
         * measure (such as centimeters).
         */
        val distanceToScreen = 60.0
        /**
         * The height of the screen area in the same distance measure (such as
         * centimeters).
         */
        val screenHeight = 32.5
        /**
         * The vertical resolution of the screen in pixels.
         */
        val screenHeightPx = 1200
    }

}

actual object GlslBase {
    actual val manager: GlslManager by lazy { JvmGlslManager() }
}

class JvmGlslManager : GlslManager {
    private val window: Long

    /**
     * This is initialization stuff that's required on the main thread.
     */
    init {
        glfwSetErrorCallback(GLFWErrorCallback.createPrint(System.err))
        if (!glfwInit())
            throw IllegalStateException("Unable to initialize GLFW")

        glfwDefaultWindowHints()
        glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4)
        glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 1)
        glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE)
        glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GLFW_TRUE)
        glfwWindowHint(GLFW_VISIBLE, GLFW_FALSE)
        glfwWindowHint(GLFW_RESIZABLE, GLFW_TRUE)
        glfwWindowHint(GLFW_SAMPLES, 8)
        glfwWindowHint(GLFW_CONTEXT_RELEASE_BEHAVIOR, GLFW_RELEASE_BEHAVIOR_NONE)

        window = glfwCreateWindow(300, 300, "Hello shaders!", NULL, NULL)
        if (window == NULL)
            throw RuntimeException("Failed to create the GLFW window")

        glfwPollEvents() // Get the event loop warmed up.
    }

    override fun getRenderer(program: String): GlslRenderer = JvmGlslRenderer(program, window)
}

class JvmGlslRenderer(private val shader: String, private var window: Long) : GlslRenderer {
    internal lateinit var keyCallback: GLFWKeyCallback
    internal lateinit var fbCallback: GLFWFramebufferSizeCallback

    internal var width = 300
    internal var height = 300
    internal var lock = Any()
    internal var destroyed: Boolean = false

    internal var viewProjMatrix = Matrix4f()
    internal var fb = BufferUtils.createFloatBuffer(16)

    val surfaces = mutableListOf<Surface>()
    var pixelCount: Int = 0

    private var vaoId: Int = 0
    private var vboId: Int = 0
    private var vertexCount: Int = 0

    private var program: Int = 0

    private var uvCoordsLocation: Int = 0
    private var pixelCountLocation: Int = 0
    private var matLocation: Int = 0
    private var resolutionLocation: Int = 0
    private var timeLocation: Int = 0

    val uvCoordTextureIndex = 0

    val outputToFrameBuffer = true

    lateinit var instance: Instance

    init {
        glfwSetKeyCallback(window, object : GLFWKeyCallback() {
            override fun invoke(window: Long, key: Int, scancode: Int, action: Int, mods: Int) {
                if (key == GLFW_KEY_ESCAPE && action == GLFW_RELEASE)
                    glfwSetWindowShouldClose(window, true)
            }
        })

        glfwSetFramebufferSizeCallback(window, object : GLFWFramebufferSizeCallback() {
            override fun invoke(window: Long, w: Int, h: Int) {
                if (outputToFrameBuffer) {
                    width = pixelCount
                    height = 1
                } else if (w > 0 && h > 0) {
                    width = w
                    height = h
                }
            }
        })

        val vidmode = glfwGetVideoMode(glfwGetPrimaryMonitor())
        glfwSetWindowPos(window, (vidmode!!.width() - width) / 2, (vidmode.height() - height) / 2)
        if (!outputToFrameBuffer) {
            glfwShowWindow(window)
        }

        val framebufferSize = BufferUtils.createIntBuffer(2)
        nglfwGetFramebufferSize(window, memAddress(framebufferSize), memAddress(framebufferSize) + 4)
        width = framebufferSize.get(0)
        height = framebufferSize.get(1)

        withGlContext(window) {
            // Create a view-projection matrix
//            val fovy = Math.atan(
//                (ViewSettings.screenHeight * height / ViewSettings.screenHeightPx)
//                        / ViewSettings.distanceToScreen
//            ).toFloat()
//            viewProjMatrix.setPerspective(fovy, 1.0f, 0.01f, 100.0f)
//                .lookAt(
//                    0.0f, 0.0f, 1.0f,
//                    0.0f, 0.0f, 0.0f,
//                    0.0f, 1.0f, 0.0f
//                )

            viewProjMatrix.setOrtho(0f, pixelCount.toFloat(), 0f, 1f, 0f, 10f)

            initQuad()

            gl { glClearColor(0f, 0f, 0f, 1f) }
            gl { glEnable(GL_DEPTH_TEST) }
            gl { glEnable(GL_CULL_FACE) }

            program = createShaderProgram()

            // Obtain uniform location
            uvCoordsLocation = gl { glGetUniformLocation(program, "sm_uvCoords") }
            pixelCountLocation = gl { glGetUniformLocation(program, "sm_pixelCount") }
            matLocation = gl { glGetUniformLocation(program, "viewProjMatrix") }
            resolutionLocation = gl { glGetUniformLocation(program, "resolution") }
            timeLocation = gl { glGetUniformLocation(program, "time") }

            instance = Instance(1, FloatArray(2))
        }
    }

    private fun initQuad() {
        // OpenGL expects vertices to be defined counter clockwise by default
        val vertices = floatArrayOf(
            // Left bottom triangle
            -0.5f, 0.5f, 0f, -0.5f, -0.5f, 0f, 0.5f, -0.5f, 0f,
            // Right top triangle
            0.5f, -0.5f, 0f, 0.5f, 0.5f, 0f, -0.5f, 0.5f, 0f
        )
        // Sending data to OpenGL requires the usage of (flipped) byte buffers
        val verticesBuffer = BufferUtils.createFloatBuffer(vertices.size)
        verticesBuffer.put(vertices)
        verticesBuffer.flip()

        vertexCount = 6

        // Create a new Vertex Array Object in memory and select it (bind)
        // A VAO can have up to 16 attributes (VBO's) assigned to it by default
        vaoId = gl { GL30.glGenVertexArrays() }
        gl { GL30.glBindVertexArray(vaoId) }

        // Create a new Vertex Buffer Object in memory and select it (bind)
        // A VBO is a collection of Vectors which in this case resemble the location of each vertex.
        vboId = gl { GL15.glGenBuffers() }
        gl { GL15.glBindBuffer(GL15.GL_ARRAY_BUFFER, vboId) }
        gl { GL15.glBufferData(GL15.GL_ARRAY_BUFFER, verticesBuffer, GL15.GL_STATIC_DRAW) }
        // Put the VBO in the attributes list at index 0
        gl { GL20.glVertexAttribPointer(0, 3, GL11.GL_FLOAT, false, 0, 0) }
        // Deselect (bind to 0) the VBO
        gl { GL15.glBindBuffer(GL15.GL_ARRAY_BUFFER, 0) }

        // Deselect (bind to 0) the VAO
        gl { GL30.glBindVertexArray(0) }
    }

    override fun addSurface(surface: Surface): Pixels {
        if (surface is IdentifiedSurface) {
            val pixelVertices = surface.pixelVertices
            if (pixelVertices != null) {
                val oldUvCoords = instance.uvCoords
                instance.release()

                val oldPixelCount = pixelCount
                val newPixelCount = oldPixelCount + surface.pixelCount

                val newUvCoords = FloatArray(newPixelCount * 2)
                oldUvCoords.copyInto(newUvCoords)
                for (i in 0 until surface.pixelCount) {
                    newUvCoords[i * 2] = pixelVertices[i].x     // u
                    newUvCoords[i * 2 + 1] = pixelVertices[i].y // v
                }

                withGlContext(window) {
                    instance = Instance(newPixelCount, newUvCoords)
                }
                pixelCount = newPixelCount

                surfaces.add(surface)
            }

            return object : Pixels {
                override val size: Int get() = surface.pixelCount

                override fun get(i: Int): Color {
                    val offset = i * 4
                    return Color(
                        instance.pixelBuffer[offset + 3], // A
                        instance.pixelBuffer[offset],     // R
                        instance.pixelBuffer[offset + 1], // G
                        instance.pixelBuffer[offset + 2]  // B
                    )
                }

                override fun set(i: Int, color: Color): Unit = TODO("set not implemented")

                override fun set(colors: Array<Color>): Unit = TODO("set not implemented")
            }
        } else {
            return object : Pixels {
                override val size: Int get() = 0

                override fun get(i: Int): Color = TODO("get not implemented")

                override fun set(i: Int, color: Color): Unit = TODO("set not implemented")

                override fun set(colors: Array<Color>): Unit = TODO("set not implemented")
            }
        }
    }

    fun runStandalone() {
        try {
            /* Spawn a new thread which to make the OpenGL context current in and which does the rendering. */
            Thread(Runnable {
                withGlContext(window) {
                    glfwSwapInterval(0)

                    while (!destroyed) {
                        render()

                        synchronized(lock) {
                            if (!destroyed) {
                                gl { glfwSwapBuffers(window) }
                            }
                        }

                    }

                    freeQuad()
                    instance.release()
                }
            }).start()

            /* Process window messages in the main thread */
            while (!glfwWindowShouldClose(window)) {
                glfwWaitEvents()
            }

            synchronized(lock) {
                destroyed = true
                glfwDestroyWindow(window)
            }
            keyCallback.free()
            fbCallback.free()
        } finally {
            glfwTerminate()
        }
    }

    internal fun renderQuad() {
//        gl { glPolygonMode(GL_FRONT, GL_FILL) }

        // Bind to the VAO that has all the information about the quad vertices
        gl { GL30.glBindVertexArray(vaoId) }
        gl { GL20.glEnableVertexAttribArray(0) }

        // Draw the vertices
        gl { GL11.glDrawArrays(GL11.GL_TRIANGLES, 0, vertexCount) }

        // Put everything back to default (deselect)
        gl { GL20.glDisableVertexAttribArray(0) }
        gl { GL30.glBindVertexArray(0) }
    }

    private fun freeQuad() {
        // Disable the VBO index from the VAO attributes list
        gl { GL20.glDisableVertexAttribArray(0) }

        // Delete the VBO
        gl { GL15.glBindBuffer(GL15.GL_ARRAY_BUFFER, 0) }
        gl { GL15.glDeleteBuffers(vboId) }

        // Delete the VAO
        gl { GL30.glBindVertexArray(0) }
        gl { GL30.glDeleteVertexArrays(vaoId) }
    }

    private fun createShaderProgram(): Int {
        // Create a simple shader program
        val program = gl { glCreateProgram() }
        val vs = gl { glCreateShader(GL_VERTEX_SHADER) }
        glShaderSource(
            vs,
            """
    #version 330 core
    
    in vec4 Vertex;
    uniform mat4 viewProjMatrix;
    
    void main(void) {
        gl_Position = viewProjMatrix * Vertex;
    }
    """
        )
        compileShader(vs)

        gl { glAttachShader(program, vs) }
        val fs = gl { glCreateShader(GL_FRAGMENT_SHADER) }
        val src = """
    #version 330
    uniform samplerBuffer sm_uvCoords;
    uniform float sm_pixelCount;
    
    out vec4 sm_fragColor;
    
    ${shader
            .replace(Regex("void main\\s*\\(\\s*void\\s*\\)"), "void sm_main(vec2 sm_pixelCoord)")
            .replace("gl_FragCoord", "sm_pixelCoord")
            .replace("gl_FragColor", "sm_fragColor")
        }
                
    void main(void) {
        int pixI = int(gl_FragCoord.x - 0.5);
        
        vec2 pixelCoord = vec2(
            texelFetch(sm_uvCoords, pixI * 2).r + .5,    // u
            texelFetch(sm_uvCoords, pixI * 2 + 1).r + .5 // v
        );
    
        sm_main(pixelCoord * resolution);
    }
    """
        println(src)
        gl { glShaderSource(fs, src) }

        compileShader(fs)

        gl { glAttachShader(program, fs) }
        gl { glLinkProgram(program) }
        if (GL20.glGetProgrami(program, GL20.GL_LINK_STATUS) == GL11.GL_FALSE) {
            throw RuntimeException("ProgramInfoLog: ${glGetProgramInfoLog(program)}")
        }

        gl { glUseProgram(program) }
        return program
    }

    override fun draw() {
//        glfwSwapInterval(0)
//        glfwPollEvents()

        val elaspedNs = timeSync {
            withGlContext(window) {
                println("Render! ${this@JvmGlslRenderer}")
                instance.bindFramebuffer()
                render()

                instance.pixelBuffer.position(0)
                GL11.glReadPixels(
                    0, 0, pixelCount, 1,
                    GL_RGBA, GL_UNSIGNED_INT_8_8_8_8, instance.pixelBuffer
                )
            }
        }

        println("Render took ${round(elaspedNs / 10000f)/100}ms")
    }

    private fun render() {
        val thisTime = System.nanoTime()

        if (outputToFrameBuffer) {
            gl { glViewport(0, 0, pixelCount, 1) }
        } else {
            gl { glViewport(0, 0, width, height) }
        }
        gl { glClear(GL_COLOR_BUFFER_BIT or GL_DEPTH_BUFFER_BIT) }

        // Upload the matrix stored in the FloatBuffer to the
        // shader uniform.
        gl { glUniformMatrix4fv(matLocation, false, viewProjMatrix.get(fb)) }
        gl { glUniform1f(pixelCountLocation, pixelCount.toFloat()) }
        gl { glUniform2f(resolutionLocation, 1000f, 1000f) }
        gl { glUniform1f(timeLocation, thisTime / 1E9f) }

        instance.bindUvCoordTexture(uvCoordTextureIndex, uvCoordsLocation)

        renderQuad()

        val programLog = gl { glGetProgramInfoLog(program) }
        if (programLog.isNotEmpty()) println("ProgramInfoLog: $programLog")
    }

    companion object {
        fun <T> gl(fn: () -> T): T {
            val result = fn.invoke()
            checkForGlError()
            return result
        }

        fun withGlContext(context: Long, fn: () -> Unit) {
            glfwMakeContextCurrent(context)
            GL.createCapabilities()

            try {
                fn()
            } finally {
                glfwMakeContextCurrent(0)
            }
        }

        fun checkForGlError() {
            while (true) {
                val error = GL11.glGetError()
                val code = when (error) {
                    GL_INVALID_ENUM -> "GL_INVALID_ENUM"
                    GL_INVALID_VALUE -> "GL_INVALID_VALUE"
                    GL_INVALID_OPERATION -> "GL_INVALID_OPERATION"
                    GL_STACK_OVERFLOW -> "GL_STACK_OVERFLOW"
                    GL_STACK_UNDERFLOW -> "GL_STACK_UNDERFLOW"
                    GL_OUT_OF_MEMORY -> "GL_OUT_OF_MEMORY"
                    else -> "unknown error $error"
                }
                if (error != 0) throw RuntimeException("OpenGL Error: $code") else return
            }
        }
    }

    private fun compileShader(shader: Int) {
        gl { glCompileShader(shader) }

        if (GL20.glGetShaderi(shader, GL20.GL_COMPILE_STATUS) == GL11.GL_FALSE) {
            throw RuntimeException("Failed to compile shader: ${glGetShaderInfoLog(shader)}")
        }
    }

    inner class Instance(val pixelCount: Int, val uvCoords: FloatArray) {
        private var uvCoordBuffer: Int = gl { glGenBuffers() }
        private var uvCoordTexture: Int = gl { GL11.glGenTextures() }
        private var frameBuffer: Int = gl { GL30.glGenFramebuffers() }
        private var renderBuffer: Int = gl { GL30.glGenRenderbuffers() }
        val pixelBuffer: ByteBuffer = ByteBuffer.allocateDirect(pixelCount * 4)

        fun bindFramebuffer() {
            gl { GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, frameBuffer) }

            gl { GL30.glBindRenderbuffer(GL30.GL_RENDERBUFFER, renderBuffer) }
            gl { GL30.glRenderbufferStorage(GL30.GL_RENDERBUFFER, GL_RGBA8, max(1, pixelCount), 1) }
            gl {
                GL30.glFramebufferRenderbuffer(
                    GL30.GL_READ_FRAMEBUFFER, GL30.GL_COLOR_ATTACHMENT0,
                    GL30.GL_RENDERBUFFER, renderBuffer
                )
            }

            val status = glCheckFramebufferStatus(GL_READ_FRAMEBUFFER)
            if (status != GL_FRAMEBUFFER_COMPLETE) {
                throw java.lang.RuntimeException("FrameBuffer huh? $status")
            }
        }

        fun bindUvCoordTexture(textureIndex: Int, uvCoordsLocation: Int) {
            gl { glBindBuffer(GL_TEXTURE_BUFFER, uvCoordBuffer) }
            gl { GL15.glBufferData(GL_TEXTURE_BUFFER, uvCoords, GL15.GL_STATIC_READ) }

            gl { glActiveTexture(GL_TEXTURE0 + textureIndex) }
            gl { glBindTexture(GL_TEXTURE_1D, uvCoordTexture) }
            gl { glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_BASE_LEVEL, 0) }
            gl { glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MAX_LEVEL, 0) }
            gl { glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) }
            gl { glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) }
            gl { glTexBuffer(GL_TEXTURE_BUFFER, GL_R32F, uvCoordBuffer) }
            gl { glUniform1i(uvCoordsLocation, textureIndex) }

            gl { glBindBuffer(GL_TEXTURE_BUFFER, 0) }
        }

        fun release() {
            if (renderBuffer != 0) {
                gl { GL30.glBindRenderbuffer(GL30.GL_RENDERBUFFER, 0) }

                gl { GL30.glDeleteRenderbuffers(renderBuffer) }
                renderBuffer = 0
            }

            gl { GL15.glDeleteBuffers(uvCoordBuffer) }
            gl { GL15.glDeleteTextures(uvCoordTexture) }
        }
    }

}