package com.performancecoachlab.posedetection.encoding

import android.graphics.Bitmap
import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.MediaMuxer
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.asAndroidBitmap
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
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 = ((width + 1) / 2) * ((height + 1) / 2) * 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 raw = frame.asAndroidBitmap()
            val bitmap = if (raw.width != width || raw.height != height) {
                Bitmap.createScaledBitmap(raw, width, height, true)
            } else {
                raw
            }
            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()
    }
}
