package com.example.vibeeffects.genai

import android.content.Context
import android.opengl.GLES20
import android.util.Log
import androidx.media3.common.C
import androidx.media3.common.VideoFrameProcessingException
import androidx.media3.common.util.GlProgram
import androidx.media3.common.util.GlUtil
import androidx.media3.common.util.GlUtil.GlException
import androidx.media3.common.util.Size
import androidx.media3.common.util.UnstableApi
import androidx.media3.common.util.Util
import androidx.media3.effect.BaseGlShaderProgram
import com.example.vibeeffects.utils.Uniform
import org.json.JSONArray
import org.json.JSONException
import java.io.IOException

/**
 * A {@link GlShaderProgram} custom shader program, likely generated by GenAI, that adds specific
 * logic to an otherwise generic template.  The template is as follows:
 *   #version 100
 *   // ES 2 fragment shader that samples from a (non-external) texture with uTexSampler,
 *   precision mediump float;
 *   uniform sampler2D uTexSampler;
 *   uniform float uTimeSeconds;
 *   varying vec2 vTexSamplingCoord;
 *   void main() {
 *     // Generated code goes here
 *   }
 */
@UnstableApi
open class GenAICustomShaderProgram(context: Context, fragmentShaderGlsl: String, val uniforms: List<Uniform<*>>, useHdr: Boolean) : BaseGlShaderProgram(useHdr, 1) {
    private val glProgram: GlProgram

    private var useTime = false
    private var useResolution = false

    private var inputWidth: Int = C.LENGTH_UNSET
    private var inputHeight: Int = C.LENGTH_UNSET

    init {
        try {
            val vertexShaderGlsl = Util.loadAsset(context, VERTEX_SHADER_PATH)
            glProgram = GlProgram(vertexShaderGlsl, fragmentShaderGlsl)
        } catch (e: IOException) {
            throw VideoFrameProcessingException(e)
        } catch (e: GlUtil.GlException) {
            throw VideoFrameProcessingException(e)
        }

        setUseOptionalInputs(fragmentShaderGlsl)

        // Draw the frame on the entire normalized device coordinate space, from -1 to 1, for x and y.
        glProgram.setBufferAttribute(
            "aFramePosition",
            GlUtil.getNormalizedCoordinateBounds(),
            GlUtil.HOMOGENEOUS_COORDINATE_VECTOR_SIZE
        )
    }

    protected fun countOccurrences(text: String, target: String): Int {
        var count = 0
        var index = 0
        while ((text.indexOf(target, index).also { index = it }) != -1) {
            count++
            index += target.length
        }
        return count
    }

    protected fun setUseOptionalInputs(glslCode: String) {
        val timeCount = countOccurrences(glslCode, U_TIME_SECONDS)
        useTime = timeCount > 1
        val resolutionCount = countOccurrences(glslCode, U_RESOLUTION)
        useResolution = resolutionCount > 1
    }

    override fun configure(
        inputWidth: Int,
        inputHeight: Int
    ): Size {
        this.inputWidth = inputWidth
        this.inputHeight = inputHeight
        return Size(inputWidth, inputHeight)
    }

    override fun drawFrame(inputTexId: Int, presentationTimeUs: Long) {
        try {
            glProgram.use()
            glProgram.setSamplerTexIdUniform("uTexSampler", inputTexId,  /* texUnitIndex= */0)
            setUniforms(glProgram, presentationTimeUs)
            glProgram.bindAttributesAndUniforms()
            // The four-vertex triangle strip forms a quad.
            GLES20.glDrawArrays(GLES20.GL_TRIANGLE_STRIP,  /* first= */0,  /* count= */4)
        } catch (e: GlException) {
            throw VideoFrameProcessingException(e, presentationTimeUs)
        }
    }

    protected fun setUniforms(glProgram: GlProgram, presentationTimeUs: Long) {
        if (useTime) {
            glProgram.setFloatUniform("uTimeSeconds", presentationTimeUs / 1000000f)
        }
        if (useResolution) {
            glProgram.setFloatsUniform(
                "uResolution", floatArrayOf(inputWidth.toFloat(), inputHeight.toFloat())
            )
        }

        for (i in uniforms.indices) {
            val uniform = uniforms[i]
            // In a real implementation, you would use the current value of the uniform,
            // not just the default. This is where you would link to UI controls.
            // For this parser, we just use the default value from the JSON.
            if (uniform.getValue() != null) {
                setUniformValue(glProgram, uniform.name, uniform.type, uniform.getValue())
            }
        }
    }

    protected fun setUniformValue(glProgram: GlProgram, name: String, type: String, value: Any?) {
        try {
            when (type) {
                "float" -> if (value is Number) {
                    glProgram.setFloatUniform(name, value.toFloat())
                }

                "int" -> if (value is Number) {
                    glProgram.setIntUniform(name, value.toInt())
                }

                "vec2", "vec3", "vec4" -> if (value is JSONArray) {
                    val array = value
                    val vec = FloatArray(array.length())
                    var i = 0
                    while (i < array.length()) {
                        vec[i] = array.getDouble(i).toFloat()
                        i++
                    }
                    glProgram.setFloatsUniform(name, vec)
                }
            }
        } catch (e: JSONException) {
            Log.e(TAG, "Failed to set uniform value for $type", e)
        }
    }

    @Throws(VideoFrameProcessingException::class)
    override fun release() {
        super.release()
        try {
            glProgram.delete()
        } catch (e: GlException) {
            throw VideoFrameProcessingException(e)
        }
    }

    companion object {
        const val TAG = "GenAICustomShaderProgram"
        const val VERTEX_SHADER_PATH = "vertex_shader_copy_es2.glsl"

        const val U_TIME_SECONDS = "uTimeSeconds"
        const val U_RESOLUTION = "uResolution"
    }
}