package com.performancecoachlab.posedetection.encoding

import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.asAndroidBitmap
import android.graphics.Bitmap
import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.MediaMuxer
import com.performancecoachlab.posedetection.recording.InputFrame
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.File
import java.io.FileOutputStream
import java.nio.ByteBuffer
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
import java.util.concurrent.ConcurrentLinkedQueue

actual fun createVideoBuilder(
    outputPath: String,
    fps: Int,
    width: Int,
    height: Int
): VideoBuilder = AndroidVideoBuilder(outputPath, fps, width, height)

class AndroidVideoBuilder(
    private val outputPath: String,
    private val fps: Int,
    private var width: Int,
    private var height: Int
) : VideoBuilder {
    private var initialised = false
    private val encoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
    private val format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height).apply {
        setInteger(MediaFormat.KEY_BIT_RATE, width * height * 4)
        setInteger(MediaFormat.KEY_FRAME_RATE, fps)
        setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1)
        // Use a specific format instead of Flexible
        setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar)
    }
    private lateinit var muxer: MediaMuxer
    private var trackIndex = -1
    private var muxerStarted = false
    private val bufferInfo = MediaCodec.BufferInfo()
    private val timeoutUs = 10000L
    private var frameIndex = 0

    private val frameQueue = ConcurrentLinkedQueue<Pair<ImageBitmap, Long>>()
    private val mutex = Mutex()
    private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
    private var processing = false
    private var isReleased = false

    private var endOfStreamSignaled = false
    private var finalizeContinuation: (()->Unit)? = null

    private var lastTimestampUs: Long = 0L

    override suspend fun addFrame(frame: ImageBitmap, timestampms: Long) {
        if (isReleased) return
        frameQueue.add(Pair(frame, timestampms*1000L))
        processQueue()
    }

    // Only process output in one place (the processing coroutine)
    private fun processQueue() {
        if (processing || isReleased) return
        processing = true
        scope.launch {
            mutex.withLock {
                while (frameQueue.isNotEmpty() || endOfStreamSignaled) {
                    if (isReleased) break
                    val pair = frameQueue.poll()
                    if (pair != null) {
                        val (nextFrame, ts) = pair
                        try {
                            processFrame(nextFrame, ts)
                        } catch (e: IllegalStateException) {
                            isReleased = true
                            frameQueue.clear()
                            break
                        }
                    } else if (endOfStreamSignaled) {
                        // Signal EOS to encoder
                        val inputBufferIndex = encoder.dequeueInputBuffer(timeoutUs)
                        if (inputBufferIndex >= 0) {
                            encoder.queueInputBuffer(
                                inputBufferIndex, 0, 0,
                                (if (frameIndex > 0) lastTimestampUs else 0L),
                                MediaCodec.BUFFER_FLAG_END_OF_STREAM
                            )
                            endOfStreamSignaled = false
                        }
                    }
                    // Drain output after each input
                    drainEncoderOutput()
                }
                processing = false
                finalizeContinuation?.invoke()
            }
        }
    }
    private fun convertBitmapToYuv420(bitmap: Bitmap): ByteArray {
        val width = bitmap.width
        val height = bitmap.height

        // NV12 format: y plane followed by interleaved uv plane
        val ySize = width * height
        val uvSize = ySize / 2
        val yuvSize = ySize + uvSize
        val yuv = ByteArray(yuvSize)

        val argb = IntArray(width * height)
        bitmap.getPixels(argb, 0, width, 0, 0, width, height)

        // Fill Y plane
        var yIndex = 0
        for (j in 0 until height) {
            for (i in 0 until width) {
                val pixel = argb[j * width + i]
                val r = (pixel shr 16) and 0xff
                val g = (pixel shr 8) and 0xff
                val b = pixel and 0xff

                // Y conversion
                val y = (0.299 * r + 0.587 * g + 0.114 * b).toInt().coerceIn(0, 255)
                yuv[yIndex++] = y.toByte()
            }
        }

        // Fill UV plane (NV12 - interleaved U and V at half resolution)
        var uvIndex = ySize
        for (j in 0 until height step 2) {
            for (i in 0 until width step 2) {
                var sumR = 0
                var sumG = 0
                var sumB = 0
                var count = 0

                // Get average color from 2x2 block
                for (y in 0 until 2) {
                    for (x in 0 until 2) {
                        val jj = j + y
                        val ii = i + x
                        if (jj < height && ii < width) {
                            val pixel = argb[jj * width + ii]
                            sumR += (pixel shr 16) and 0xff
                            sumG += (pixel shr 8) and 0xff
                            sumB += pixel and 0xff
                            count++
                        }
                    }
                }

                if (count > 0) {
                    val avgR = sumR / count
                    val avgG = sumG / count
                    val avgB = sumB / count

                    // U and V conversion
                    val u = (-0.169 * avgR - 0.331 * avgG + 0.5 * avgB + 128).toInt().coerceIn(0, 255)
                    val v = (0.5 * avgR - 0.419 * avgG - 0.081 * avgB + 128).toInt().coerceIn(0, 255)

                    // NV12 has interleaved U and V values
                    yuv[uvIndex++] = u.toByte()
                    yuv[uvIndex++] = v.toByte()
                }
            }
        }

        return yuv
    }

    private suspend fun processFrame(frame: ImageBitmap, timestampUs: Long) {
        if (isReleased) return
        try {
            if (!initialised) {
                encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
                encoder.start()
                muxer = MediaMuxer(outputPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
                initialised = true
            }
            val bitmap = frame.asAndroidBitmap()
            val yuvData = convertBitmapToYuv420(bitmap)
            val inputBufferIndex = try { encoder.dequeueInputBuffer(timeoutUs) } catch (e: IllegalStateException) { isReleased = true; frameQueue.clear(); return }
            if (inputBufferIndex >= 0) {
                val inputBuffer = try { encoder.getInputBuffer(inputBufferIndex) } catch (e: IllegalStateException) { isReleased = true; frameQueue.clear(); return }
                inputBuffer?.clear()
                if (inputBuffer != null) {
                    val bufferCapacity = inputBuffer.capacity()
                    if (yuvData.size <= bufferCapacity) {
                        inputBuffer.put(yuvData)
                    } else {
                        inputBuffer.put(yuvData, 0, bufferCapacity)
                        //println("Warning: YUV data (${yuvData.size} bytes) exceeded buffer capacity ($bufferCapacity bytes)")
                    }
                }
                val dataSize = minOf(yuvData.size, inputBuffer?.capacity() ?: 0)
                try {
                    // Use the provided timestampUs for variable frame rate
                    encoder.queueInputBuffer(inputBufferIndex, 0, dataSize, timestampUs, 0)
                    lastTimestampUs = timestampUs
                } catch (e: IllegalStateException) {
                    isReleased = true; frameQueue.clear(); return
                }
            }
            frameIndex++
        } catch (e: IllegalStateException) {
            isReleased = true
            frameQueue.clear()
            return
        }
    }

    private fun drainEncoderOutput() {
        var outputDone = false
        while (!outputDone) {
            val outputBufferIndex = try { encoder.dequeueOutputBuffer(bufferInfo, timeoutUs) } catch (e: IllegalStateException) { isReleased = true; frameQueue.clear(); return }
            when {
                outputBufferIndex == MediaCodec.INFO_TRY_AGAIN_LATER -> {
                    outputDone = true
                }
                outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED -> {
                    // Set up muxer when format is ready
                    try {
                        trackIndex = muxer.addTrack(encoder.outputFormat)
                        muxer.start()
                        muxerStarted = true
                    } catch (e: IllegalStateException) {
                        isReleased = true; frameQueue.clear(); return
                    }
                }
                outputBufferIndex >= 0 -> {
                    val outputBuffer = try { encoder.getOutputBuffer(outputBufferIndex) } catch (e: IllegalStateException) { isReleased = true; frameQueue.clear(); return }
                    if (bufferInfo.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG != 0) {
                        // Config info, not actual frame data
                        bufferInfo.size = 0
                    }

                    if (bufferInfo.size > 0 && muxerStarted) {
                        try {
                            outputBuffer?.position(bufferInfo.offset)
                            outputBuffer?.limit(bufferInfo.offset + bufferInfo.size)
                            muxer.writeSampleData(trackIndex, outputBuffer!!, bufferInfo)
                        } catch (e: IllegalStateException) {
                            isReleased = true; frameQueue.clear(); return
                        }
                    }
                    try {
                        encoder.releaseOutputBuffer(outputBufferIndex, false)
                    } catch (e: IllegalStateException) {
                        isReleased = true; frameQueue.clear(); return
                    }

                    if (bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM != 0) {
                        outputDone = true
                    }
                }
            }
        }
    }

    override suspend fun finalize(): String = withContext(Dispatchers.IO) {
        if (initialised) {
            // Signal end of stream to encoder and wait for processing to finish
            val finalizeSignal = kotlinx.coroutines.CompletableDeferred<Unit>()
            finalizeContinuation = { finalizeSignal.complete(Unit) }
            endOfStreamSignaled = true
            processQueue()
            finalizeSignal.await()
            cleanup()
        }
        outputPath
    }

    override suspend fun cancel() = withContext(Dispatchers.IO) {
        cleanup()
    }

    private fun cleanup() {
        if (initialised) {
            encoder.stop()
            encoder.release()
            if (::muxer.isInitialized && muxerStarted) {
                muxer.stop()
                muxer.release()
            }
            initialised = false
        }
        isReleased = true
        frameQueue.clear()
    }
}
