tilt-shift-camera/app/src/main/java/no/naiv/tiltshift/effect/TiltShiftShader.kt

package no.naiv.tiltshift.effect

import android.content.Context
import android.opengl.GLES11Ext
import android.opengl.GLES20
import no.naiv.tiltshift.R
import java.io.BufferedReader
import java.io.InputStreamReader
import kotlin.math.cos
import kotlin.math.sin

/**
 * Manages OpenGL shader programs for the two-pass tilt-shift effect.
 *
 * Two programs:
 * - **Passthrough**: copies camera texture (external OES) to an FBO, handling the
 *   coordinate transform via vertex/texcoord setup.
 * - **Blur**: applies a directional Gaussian blur with tilt-shift mask.
 *   Used twice per frame (horizontal then vertical) via the [uBlurDirection] uniform.
 */
class TiltShiftShader(private val context: Context) {

    private companion object {
        val IDENTITY_MATRIX = floatArrayOf(
            1f, 0f, 0f, 0f,
            0f, 1f, 0f, 0f,
            0f, 0f, 1f, 0f,
            0f, 0f, 0f, 1f
        )
    }

    // --- Passthrough program (camera → FBO) ---

    private var passthroughProgramId: Int = 0

    var passthroughPositionLoc: Int = 0
        private set
    var passthroughTexCoordLoc: Int = 0
        private set
    private var passthroughTextureLoc: Int = 0
    private var passthroughTexMatrixLoc: Int = 0
    private var passthroughMirrorLoc: Int = 0

    // --- Blur program (FBO → FBO/screen) ---

    private var blurProgramId: Int = 0

    var blurPositionLoc: Int = 0
        private set
    var blurTexCoordLoc: Int = 0
        private set
    private var blurTextureLoc: Int = 0
    private var blurTexMatrixLoc: Int = 0
    private var blurMirrorLoc: Int = 0
    private var blurModeLoc: Int = 0
    private var blurPositionXLoc: Int = 0
    private var blurPositionYLoc: Int = 0
    private var blurSizeLoc: Int = 0
    private var blurAmountLoc: Int = 0
    private var blurFalloffLoc: Int = 0
    private var blurAspectRatioLoc: Int = 0
    private var blurResolutionLoc: Int = 0
    private var blurCosAngleLoc: Int = 0
    private var blurSinAngleLoc: Int = 0
    private var blurDirectionLoc: Int = 0

    /**
     * Compiles and links both shader programs.
     * Must be called from GL thread.
     */
    fun initialize() {
        val vertexSource = loadShaderSource(R.raw.tiltshift_vertex)
        val vertexShader = compileShader(GLES20.GL_VERTEX_SHADER, vertexSource)

        // Passthrough program
        val passthroughFragSource = loadShaderSource(R.raw.tiltshift_passthrough_fragment)
        val passthroughFragShader = compileShader(GLES20.GL_FRAGMENT_SHADER, passthroughFragSource)
        passthroughProgramId = linkProgram(vertexShader, passthroughFragShader)
        GLES20.glDeleteShader(passthroughFragShader)

        passthroughPositionLoc = GLES20.glGetAttribLocation(passthroughProgramId, "aPosition")
        passthroughTexCoordLoc = GLES20.glGetAttribLocation(passthroughProgramId, "aTexCoord")
        passthroughTextureLoc = GLES20.glGetUniformLocation(passthroughProgramId, "uTexture")
        passthroughTexMatrixLoc = GLES20.glGetUniformLocation(passthroughProgramId, "uTexMatrix")
        passthroughMirrorLoc = GLES20.glGetUniformLocation(passthroughProgramId, "uMirrorX")

        // Blur program
        val blurFragSource = loadShaderSource(R.raw.tiltshift_fragment)
        val blurFragShader = compileShader(GLES20.GL_FRAGMENT_SHADER, blurFragSource)
        blurProgramId = linkProgram(vertexShader, blurFragShader)
        GLES20.glDeleteShader(blurFragShader)

        blurPositionLoc = GLES20.glGetAttribLocation(blurProgramId, "aPosition")
        blurTexCoordLoc = GLES20.glGetAttribLocation(blurProgramId, "aTexCoord")
        blurTextureLoc = GLES20.glGetUniformLocation(blurProgramId, "uTexture")
        blurTexMatrixLoc = GLES20.glGetUniformLocation(blurProgramId, "uTexMatrix")
        blurMirrorLoc = GLES20.glGetUniformLocation(blurProgramId, "uMirrorX")
        blurModeLoc = GLES20.glGetUniformLocation(blurProgramId, "uMode")
        blurPositionXLoc = GLES20.glGetUniformLocation(blurProgramId, "uPositionX")
        blurPositionYLoc = GLES20.glGetUniformLocation(blurProgramId, "uPositionY")
        blurSizeLoc = GLES20.glGetUniformLocation(blurProgramId, "uSize")
        blurAmountLoc = GLES20.glGetUniformLocation(blurProgramId, "uBlurAmount")
        blurFalloffLoc = GLES20.glGetUniformLocation(blurProgramId, "uFalloff")
        blurAspectRatioLoc = GLES20.glGetUniformLocation(blurProgramId, "uAspectRatio")
        blurResolutionLoc = GLES20.glGetUniformLocation(blurProgramId, "uResolution")
        blurCosAngleLoc = GLES20.glGetUniformLocation(blurProgramId, "uCosAngle")
        blurSinAngleLoc = GLES20.glGetUniformLocation(blurProgramId, "uSinAngle")
        blurDirectionLoc = GLES20.glGetUniformLocation(blurProgramId, "uBlurDirection")

        // Vertex shader is linked into both programs and can be freed
        GLES20.glDeleteShader(vertexShader)
    }

    /**
     * Activates the passthrough program and binds the camera texture.
     *
     * @param cameraTextureId  The OES texture receiving camera frames.
     * @param texMatrix        4x4 transform from SurfaceTexture.getTransformMatrix() —
     *                         encodes sensor-to-display rotation and Y-flip.
     * @param mirrorX          true to horizontally mirror (front camera selfie view).
     */
    fun usePassthrough(cameraTextureId: Int, texMatrix: FloatArray, mirrorX: Boolean) {
        GLES20.glUseProgram(passthroughProgramId)
        GLES20.glActiveTexture(GLES20.GL_TEXTURE0)
        GLES20.glBindTexture(GLES11Ext.GL_TEXTURE_EXTERNAL_OES, cameraTextureId)
        GLES20.glUniform1i(passthroughTextureLoc, 0)
        GLES20.glUniformMatrix4fv(passthroughTexMatrixLoc, 1, false, texMatrix, 0)
        GLES20.glUniform1f(passthroughMirrorLoc, if (mirrorX) 1f else 0f)
    }

    /**
     * Activates the blur program and sets all uniforms for one blur pass.
     *
     * @param fboTextureId The FBO color attachment to sample from.
     * @param params       Current blur parameters.
     * @param width        Surface width in pixels.
     * @param height       Surface height in pixels.
     * @param dirX         Blur direction X component (1 for horizontal pass, 0 for vertical).
     * @param dirY         Blur direction Y component (0 for horizontal pass, 1 for vertical).
     */
    fun useBlurPass(
        fboTextureId: Int,
        params: BlurParameters,
        width: Int,
        height: Int,
        dirX: Float,
        dirY: Float
    ) {
        GLES20.glUseProgram(blurProgramId)

        GLES20.glActiveTexture(GLES20.GL_TEXTURE0)
        GLES20.glBindTexture(GLES20.GL_TEXTURE_2D, fboTextureId)
        GLES20.glUniform1i(blurTextureLoc, 0)

        // FBO content is already in display orientation — pass identity matrix and no mirror.
        GLES20.glUniformMatrix4fv(blurTexMatrixLoc, 1, false, IDENTITY_MATRIX, 0)
        GLES20.glUniform1f(blurMirrorLoc, 0f)

        GLES20.glUniform1i(blurModeLoc, if (params.mode == BlurMode.RADIAL) 1 else 0)
        GLES20.glUniform1f(blurPositionXLoc, params.positionX)
        GLES20.glUniform1f(blurPositionYLoc, params.positionY)
        GLES20.glUniform1f(blurSizeLoc, params.size)
        GLES20.glUniform1f(blurAmountLoc, params.blurAmount)
        GLES20.glUniform1f(blurFalloffLoc, params.falloff)
        GLES20.glUniform1f(blurAspectRatioLoc, params.aspectRatio)
        GLES20.glUniform2f(blurResolutionLoc, width.toFloat(), height.toFloat())

        // Raw screen-space angle (no camera rotation adjustment needed — FBO is already
        // in screen orientation after the passthrough pass)
        GLES20.glUniform1f(blurCosAngleLoc, cos(params.angle))
        GLES20.glUniform1f(blurSinAngleLoc, sin(params.angle))

        GLES20.glUniform2f(blurDirectionLoc, dirX, dirY)
    }

    /**
     * Releases both shader programs.
     */
    fun release() {
        if (passthroughProgramId != 0) {
            GLES20.glDeleteProgram(passthroughProgramId)
            passthroughProgramId = 0
        }
        if (blurProgramId != 0) {
            GLES20.glDeleteProgram(blurProgramId)
            blurProgramId = 0
        }
    }

    private fun linkProgram(vertexShader: Int, fragmentShader: Int): Int {
        val programId = GLES20.glCreateProgram()
        GLES20.glAttachShader(programId, vertexShader)
        GLES20.glAttachShader(programId, fragmentShader)
        GLES20.glLinkProgram(programId)

        val linkStatus = IntArray(1)
        GLES20.glGetProgramiv(programId, GLES20.GL_LINK_STATUS, linkStatus, 0)
        if (linkStatus[0] == 0) {
            val error = GLES20.glGetProgramInfoLog(programId)
            GLES20.glDeleteProgram(programId)
            throw RuntimeException("Shader program link failed: $error")
        }

        return programId
    }

    private fun loadShaderSource(resourceId: Int): String {
        val inputStream = context.resources.openRawResource(resourceId)
        val reader = BufferedReader(InputStreamReader(inputStream))
        return reader.use { it.readText() }
    }

    private fun compileShader(type: Int, source: String): Int {
        val shader = GLES20.glCreateShader(type)
        GLES20.glShaderSource(shader, source)
        GLES20.glCompileShader(shader)

        val compileStatus = IntArray(1)
        GLES20.glGetShaderiv(shader, GLES20.GL_COMPILE_STATUS, compileStatus, 0)
        if (compileStatus[0] == 0) {
            val error = GLES20.glGetShaderInfoLog(shader)
            GLES20.glDeleteShader(shader)
            val shaderType = if (type == GLES20.GL_VERTEX_SHADER) "vertex" else "fragment"
            throw RuntimeException("$shaderType shader compilation failed: $error")
        }

        return shader
    }
}