package com.nothing.camera2magic.hook

import android.graphics.SurfaceTexture
import android.opengl.*
import android.opengl.EGL14.*
import android.opengl.GLES11Ext.GL_TEXTURE_EXTERNAL_OES
import android.opengl.GLES20.*
import android.os.Handler
import android.os.HandlerThread
import android.view.Surface
import com.nothing.camera2magic.utils.Dog
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.FloatBuffer

class VirtualCameraRenderer : SurfaceTexture.OnFrameAvailableListener {
    companion object {
        private const val TAG = "VCR"
    }

    private val renderThread = HandlerThread("CameraRenderThread").apply { start() }
    private val renderHandler = Handler(renderThread.looper)

    private var eglDisplay: EGLDisplay = EGL_NO_DISPLAY
    private var eglContext: EGLContext = EGL_NO_CONTEXT
    private var eglSurface: EGLSurface = EGL_NO_SURFACE
    private var eglConfig: EGLConfig? = null

    private var surfaceTexture: SurfaceTexture? = null
    private var textureId: Int = -1
    private var programId: Int = -1
    private var uMVPMatrixHandle: Int = -1
    private var uSTMatrixHandle: Int = -1
    private var aPositionHandle: Int = -1
    private var aTextureCoordHandle: Int = -1

    private val mvpMatrix = FloatArray(16)
    private val stMatrix = FloatArray(16)

    @Volatile
    var inputSurface: Surface? = null
        private set

    private val targetSurfaces = mutableListOf<EGLSurface>()
    private val surfaceSizes = mutableMapOf<EGLSurface, Pair<Int, Int>>()
    private val surfaceFormats = mutableMapOf<EGLSurface, Int>()
    private val imageWriters = mutableMapOf<Surface, android.media.ImageWriter>()
    private var primaryWidth = 0
    private var primaryHeight = 0
    private var videoWidth = 0
    private var videoHeight = 0

    private val vertexShaderCode = """
        uniform mat4 uMVPMatrix;
        uniform mat4 uSTMatrix;
        attribute vec4 aPosition;
        attribute vec4 aTextureCoord;
        varying vec2 vTextureCoord;
        void main() {
            gl_Position = uMVPMatrix * aPosition;
            vTextureCoord = (uSTMatrix * aTextureCoord).xy;
        }
    """.trimIndent()

    private val fragmentShaderCode = """
        #extension GL_OES_EGL_image_external : require
        precision mediump float;
        varying vec2 vTextureCoord;
        uniform samplerExternalOES sTexture;
        void main() {
            gl_FragColor = texture2D(sTexture, vTextureCoord);
        }
    """.trimIndent()

    private val vertexBuffer: FloatBuffer = ByteBuffer.allocateDirect(16 * 4)
        .order(ByteOrder.nativeOrder())
        .asFloatBuffer()
        .apply {
            put(floatArrayOf(
                -1.0f, -1.0f, 0.0f, 1.0f,
                1.0f, -1.0f, 0.0f, 1.0f,
                -1.0f, 1.0f, 0.0f, 1.0f,
                1.0f, 1.0f, 0.0f, 1.0f
            ))
            position(0)
        }

    private val textureBuffer: FloatBuffer = ByteBuffer.allocateDirect(16 * 4)
        .order(ByteOrder.nativeOrder())
        .asFloatBuffer()
        .apply {
            put(floatArrayOf(
                0.0f, 0.0f, 0.0f, 1.0f,
                1.0f, 0.0f, 0.0f, 1.0f,
                0.0f, 1.0f, 0.0f, 1.0f,
                1.0f, 1.0f, 0.0f, 1.0f
            ))
            position(0)
        }

    private var pbufferSurface: EGLSurface = EGL_NO_SURFACE
    private var pixelBuffer: ByteBuffer? = null

    init {
        Matrix.setIdentityM(mvpMatrix, 0)
        Matrix.setIdentityM(stMatrix, 0)
        renderHandler.post { setupEGL() }
    }

    fun setTargetSurfaces(surfaces: Collection<Surface>, width: Int, height: Int) {
        renderHandler.post {
            Dog.i(TAG, "setTargetSurfaces: count=${surfaces.size}, ref_w=$width, ref_h=$height", true)
            
            // Clean up existing surfaces
            targetSurfaces.forEach { eglDestroySurface(eglDisplay, it) }
            targetSurfaces.clear()
            surfaceSizes.clear()
            surfaceFormats.clear()

            primaryWidth = width
            primaryHeight = height
            updateMVPMatrix()

            surfaces.forEach { surface ->
                if (surface.isValid) {
                    val (w, h, f) = NativeBridge.getSurfaceInfo(surface)
                    
                    // If it's a YUV/JPEG capture surface, we use ImageWriter instead of EGL
                    if (f == 34 || f == 35 || f == 33) {
                        Dog.i(TAG, "Creating ImageWriter for YUV/JPEG surface (format $f)", true)
                        try {
                            val writer = android.media.ImageWriter.newInstance(surface, 2)
                            imageWriters[surface] = writer
                        } catch (e: Exception) {
                            Dog.e(TAG, "Failed to create ImageWriter for format $f", e, true)
                        }
                        return@forEach
                    }

                    val surfaceAttribs = intArrayOf(EGL_NONE)
                    val eglSurface = eglCreateWindowSurface(eglDisplay, eglConfig, surface, surfaceAttribs, 0)
                    if (eglSurface != EGL_NO_SURFACE) {
                        targetSurfaces.add(eglSurface)
                        surfaceSizes[eglSurface] = Pair(w, h)
                        surfaceFormats[eglSurface] = f
                        Dog.i(TAG, "Added EGLSurface for $surface ($w x $h, format $f)", true)
                    } else {
                        val error = eglGetError()
                        Dog.e(TAG, "eglCreateWindowSurface failed for $surface: ${GLUtils.getEGLErrorString(error)}", null, true)
                    }
                }
            }
        }
    }

    private fun updateMVPMatrix() {
        if (primaryWidth <= 0 || primaryHeight <= 0 || videoWidth <= 0 || videoHeight <= 0) {
            Matrix.setIdentityM(mvpMatrix, 0)
            return
        }
        val surfaceRatio = primaryWidth.toFloat() / primaryHeight
        val videoRatio = videoWidth.toFloat() / videoHeight
        
        Matrix.setIdentityM(mvpMatrix, 0)
        if (videoRatio > surfaceRatio) {
            // Video is wider than surface - letterbox top/bottom
            Matrix.scaleM(mvpMatrix, 0, 1.0f, surfaceRatio / videoRatio, 1.0f)
        } else {
            // Video is narrower than surface - letterbox sides
            Matrix.scaleM(mvpMatrix, 0, videoRatio / surfaceRatio, 1.0f, 1.0f)
        }
        Dog.i(TAG, "MVP updated: sRatio=$surfaceRatio, vRatio=$videoRatio", true)
    }

    private fun setupEGL() {
        eglDisplay = eglGetDisplay(EGL_DEFAULT_DISPLAY)
        val version = IntArray(2)
        eglInitialize(eglDisplay, version, 0, version, 1)

        val attribList = intArrayOf(
            EGL_RED_SIZE, 8,
            EGL_GREEN_SIZE, 8,
            EGL_BLUE_SIZE, 8,
            EGL_RENDERABLE_TYPE, EGL_OPENGL_ES2_BIT,
            EGL_NONE
        )
        val configs = arrayOfNulls<EGLConfig>(1)
        val numConfigs = IntArray(1)
        eglChooseConfig(eglDisplay, attribList, 0, configs, 0, configs.size, numConfigs, 0)
        eglConfig = configs[0]

        val contextAttribs = intArrayOf(EGL_CONTEXT_CLIENT_VERSION, 2, EGL_NONE)
        eglContext = eglCreateContext(eglDisplay, eglConfig, EGL_NO_CONTEXT, contextAttribs, 0)

        val pbufferAttribs = intArrayOf(EGL_WIDTH, 1, EGL_HEIGHT, 1, EGL_NONE)
        pbufferSurface = eglCreatePbufferSurface(eglDisplay, eglConfig, pbufferAttribs, 0)
        eglMakeCurrent(eglDisplay, pbufferSurface, pbufferSurface, eglContext)

        setupShaders()
        setupTexture()

        surfaceTexture = SurfaceTexture(textureId).apply {
            setOnFrameAvailableListener(this@VirtualCameraRenderer, renderHandler)
        }
        inputSurface = Surface(surfaceTexture)
        Dog.i(TAG, "EGL Setup complete. inputSurface ready.", true)
    }

    private fun setupShaders() {
        val vertexShader = loadShader(GL_VERTEX_SHADER, vertexShaderCode)
        val fragmentShader = loadShader(GL_FRAGMENT_SHADER, fragmentShaderCode)
        programId = glCreateProgram().also {
            glAttachShader(it, vertexShader)
            glAttachShader(it, fragmentShader)
            glLinkProgram(it)
        }
        
        uMVPMatrixHandle = glGetUniformLocation(programId, "uMVPMatrix")
        uSTMatrixHandle = glGetUniformLocation(programId, "uSTMatrix")
        aPositionHandle = glGetAttribLocation(programId, "aPosition")
        aTextureCoordHandle = glGetAttribLocation(programId, "aTextureCoord")
    }

    private fun setupTexture() {
        val textures = IntArray(1)
        glGenTextures(1, textures, 0)
        textureId = textures[0]
        glBindTexture(GL_TEXTURE_EXTERNAL_OES, textureId)
        glTexParameterf(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_MIN_FILTER, GL_LINEAR.toFloat())
        glTexParameterf(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_MAG_FILTER, GL_LINEAR.toFloat())
        glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
    }

    private fun loadShader(type: Int, shaderCode: String): Int {
        return glCreateShader(type).also {
            glShaderSource(it, shaderCode)
            glCompileShader(it)
        }
    }

    override fun onFrameAvailable(st: SurfaceTexture) {
        renderHandler.post {
            if (targetSurfaces.isEmpty()) return@post
            
            st.updateTexImage()
            st.getTransformMatrix(stMatrix)
            val timestamp = st.timestamp

            val it = targetSurfaces.iterator()
            while (it.hasNext()) {
                val eglSurface = it.next()
                if (!eglMakeCurrent(eglDisplay, eglSurface, eglSurface, eglContext)) {
                    val error = eglGetError()
                    Dog.e(TAG, "eglMakeCurrent failed for surface, removing: ${GLUtils.getEGLErrorString(error)}", null, true)
                    eglDestroySurface(eglDisplay, eglSurface)
                    it.remove()
                    continue
                }
                
                val size = surfaceSizes[eglSurface] ?: Pair(primaryWidth, primaryHeight)
                glViewport(0, 0, size.first, size.second)
                
                glClearColor(0.0f, 0.0f, 0.0f, 1.0f)
                glClear(GL_COLOR_BUFFER_BIT)
                
                glUseProgram(programId)
                glUniformMatrix4fv(uMVPMatrixHandle, 1, false, mvpMatrix, 0)
                glUniformMatrix4fv(uSTMatrixHandle, 1, false, stMatrix, 0)

                glEnableVertexAttribArray(aPositionHandle)
                glVertexAttribPointer(aPositionHandle, 4, GL_FLOAT, false, 16, vertexBuffer)
                
                glEnableVertexAttribArray(aTextureCoordHandle)
                glVertexAttribPointer(aTextureCoordHandle, 4, GL_FLOAT, false, 16, textureBuffer)
                
                glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
                
                glDisableVertexAttribArray(aPositionHandle)
                glDisableVertexAttribArray(aTextureCoordHandle)
                
                if (timestamp > 0) {
                    EGLExt.eglPresentationTimeANDROID(eglDisplay, eglSurface, timestamp)
                }
                eglSwapBuffers(eglDisplay, eglSurface)
            }

            // Feed YUV ImageWriters
            if (imageWriters.isNotEmpty()) {
                feedImageWriters(st)
            }
        }
    }

    private fun feedImageWriters(st: SurfaceTexture) {
        val width = primaryWidth
        val height = primaryHeight
        if (width <= 0 || height <= 0) return

        // 1. Draw to offscreen pbuffer for reading
        eglMakeCurrent(eglDisplay, pbufferSurface, pbufferSurface, eglContext)
        // Need to resize pbuffer or use an FBO for correct sizing. 
        // For now, we assume the pbuffer is used as a state-holder and we use an FBO if needed.
        // Simplified: Read from the last drawn EGL surface if available.
        if (targetSurfaces.isNotEmpty()) {
             eglMakeCurrent(eglDisplay, targetSurfaces[0], targetSurfaces[0], eglContext)
        }

        val size = width * height * 4
        if (pixelBuffer == null || pixelBuffer!!.capacity() < size) {
            pixelBuffer = ByteBuffer.allocateDirect(size).order(ByteOrder.nativeOrder())
        }
        pixelBuffer!!.rewind()
        glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE, pixelBuffer)

        imageWriters.forEach { (surface, writer) ->
            try {
                val image = writer.dequeueInputImage()
                val planes = image.planes
                
                // Fast RGB to YUV conversion (simplified I420)
                val yBuffer = planes[0].buffer
                val uBuffer = planes[1].buffer
                val vBuffer = planes[2].buffer
                
                val yStride = planes[0].rowStride
                val uvStride = planes[1].rowStride
                val uvPixelStride = planes[1].pixelStride
                
                pixelBuffer!!.rewind()
                for (y in 0 until height) {
                    for (x in 0 until width) {
                        val r = pixelBuffer!!.get().toInt() and 0xFF
                        val g = pixelBuffer!!.get().toInt() and 0xFF
                        val b = pixelBuffer!!.get().toInt() and 0xFF
                        pixelBuffer!!.get() // skip A
                        
                        // Y = 0.299R + 0.587G + 0.114B
                        val yVal = ((66 * r + 129 * g + 25 * b + 128) shr 8) + 16
                        yBuffer.put(y * yStride + x, yVal.toByte())
                        
                        if (y % 2 == 0 && x % 2 == 0) {
                            // U = -0.169R - 0.331G + 0.500B
                            val uVal = ((-38 * r - 74 * g + 112 * b + 128) shr 8) + 128
                            // V = 0.500R - 0.419G - 0.081B
                            val vVal = ((112 * r - 94 * g - 18 * b + 128) shr 8) + 128
                            
                            uBuffer.put((y / 2) * uvStride + (x / 2) * uvPixelStride, uVal.toByte())
                            vBuffer.put((y / 2) * uvStride + (x / 2) * uvPixelStride, vVal.toByte())
                        }
                    }
                }
                writer.queueInputImage(image)
            } catch (e: Exception) {
                // Ignore if writer is full or surface is gone
            }
        }
    }

    fun setVideoSize(width: Int, height: Int) {
        renderHandler.post {
            Dog.i(TAG, "setVideoSize: $width x $height", true)
            videoWidth = width
            videoHeight = height
            updateMVPMatrix()
        }
    }

    fun release() {
        renderHandler.post {
            targetSurfaces.forEach { eglDestroySurface(eglDisplay, it) }
            targetSurfaces.clear()
            surfaceSizes.clear()
            if (eglContext != EGL_NO_CONTEXT) {
                eglDestroyContext(eglDisplay, eglContext)
                eglContext = EGL_NO_CONTEXT
            }
            if (eglDisplay != EGL_NO_DISPLAY) {
                eglTerminate(eglDisplay)
                eglDisplay = EGL_NO_DISPLAY
            }
            inputSurface?.release()
            surfaceTexture?.release()
            renderThread.quit()
        }
    }
}
