package com.performancecoachlab.posedetection.recording

import android.app.Application
import android.content.Context
import android.graphics.Bitmap
import android.graphics.Matrix
import android.media.Image
import android.media.MediaCodec
import android.media.MediaExtractor
import android.media.MediaFormat
import androidx.core.net.toUri
import co.touchlab.kermit.Logger
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.flowOn
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlin.coroutines.cancellation.CancellationException

/**
 * Sequential video decoder using MediaCodec. Opens the video once and decodes
 * frames in presentation order. Much faster than MediaMetadataRetriever which
 * re-opens/re-seeks per frame.
 *
 * Frames are requested in ascending timestamp order (as produced by
 * [listVideoFrameTimestamps]). The decoder advances until it reaches a frame
 * whose PTS is >= the requested timestamp, then returns that frame.
 */
private class VideoFrameDecoder(private val videoPath: String) {
    private var extractor: MediaExtractor? = null
    private var codec: MediaCodec? = null
    private var initialised = false
    private var inputEos = false
    private var outputEos = false
    private var rotationDegrees = 0

    // The most recently decoded frame and its PTS (ms).
    private var lastBitmap: Bitmap? = null
    private var lastPtsMs: Long = -1


    private fun init() {
        if (initialised) return
        val ctx = VideoExtractionContext.get()
        val uri = videoPath.toUri()
        val ext = MediaExtractor()
        ext.setDataSource(ctx, uri, null)

        for (i in 0 until ext.trackCount) {
            val format = ext.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME) ?: continue
            if (mime.startsWith("video/")) {
                ext.selectTrack(i)
                rotationDegrees = if (format.containsKey(MediaFormat.KEY_ROTATION)) {
                    format.getInteger(MediaFormat.KEY_ROTATION)
                } else 0
                val decoder = MediaCodec.createDecoderByType(mime)
                decoder.configure(format, null, null, 0)
                decoder.start()
                codec = decoder
                extractor = ext
                initialised = true
                return
            }
        }
        ext.release()
        Logger.w { "VideoFrameDecoder: no video track in $videoPath" }
    }

    /**
     * Decode forward until we produce a frame at or past [targetMs].
     * Returns the decoded bitmap tagged with [targetMs] as the InputFrame timestamp
     * (so the caller's relative-timestamp math stays correct).
     */
    fun decodeUpTo(targetMs: Long): Bitmap? {
        if (!initialised) init()
        val decoder = codec ?: return null
        val ext = extractor ?: return null

        // Already past this timestamp — return last frame.
        if (lastPtsMs >= targetMs && lastBitmap != null) {
            return lastBitmap
        }

        val info = MediaCodec.BufferInfo()

        try {
            while (!outputEos) {
                // Feed as many input packets as possible to keep the decoder pipeline full.
                if (!inputEos) {
                    while (true) {
                        val inIdx = decoder.dequeueInputBuffer(10_000L)
                        if (inIdx < 0) break
                        val buf = decoder.getInputBuffer(inIdx)!!
                        val size = ext.readSampleData(buf, 0)
                        if (size < 0) {
                            decoder.queueInputBuffer(
                                inIdx, 0, 0, 0, MediaCodec.BUFFER_FLAG_END_OF_STREAM
                            )
                            inputEos = true
                            break
                        } else {
                            decoder.queueInputBuffer(inIdx, 0, size, ext.sampleTime, 0)
                            ext.advance()
                        }
                    }
                }

                // Drain output.
                val outIdx = decoder.dequeueOutputBuffer(info, 10_000L)
                if (outIdx >= 0) {
                    val ptsMs = info.presentationTimeUs / 1000L

                    if (info.size > 0) {
                        val image = decoder.getOutputImage(outIdx)
                        if (image != null) {
                            val raw = yuvImageToBitmap(image)
                            image.close()
                            lastBitmap = applyRotation(raw)
                            lastPtsMs = ptsMs
                        }
                    }

                    val eos = info.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM != 0
                    decoder.releaseOutputBuffer(outIdx, false)

                    if (eos) {
                        outputEos = true
                        return lastBitmap
                    }

                    // We've reached or passed the target — return this frame.
                    if (ptsMs >= targetMs) {
                        return lastBitmap
                    }
                }
            }
        } catch (e: Exception) {
            Logger.w { "decodeUpTo failed at targetMs=$targetMs: ${e.message}" }
            release()
            throw e
        }

        return lastBitmap
    }

    fun release() {
        try { codec?.stop() } catch (_: Exception) {}
        try { codec?.release() } catch (_: Exception) {}
        codec = null
        try { extractor?.release() } catch (_: Exception) {}
        extractor = null
        lastBitmap = null
        lastPtsMs = -1
        initialised = false
        inputEos = false
        outputEos = false
    }

    private fun applyRotation(bitmap: Bitmap): Bitmap {
        if (rotationDegrees == 0) return bitmap
        val matrix = Matrix()
        matrix.postRotate(rotationDegrees.toFloat())
        return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
    }

    /**
     * Fast YUV Image to ARGB Bitmap conversion using bulk buffer copies
     * and fixed-point integer arithmetic (no floating-point per pixel).
     */
    private fun yuvImageToBitmap(image: Image): Bitmap {
        val w = image.width
        val h = image.height
        val yPlane = image.planes[0]
        val uPlane = image.planes[1]
        val vPlane = image.planes[2]

        val yRowStride = yPlane.rowStride
        val uvRowStride = uPlane.rowStride
        val uvPixelStride = uPlane.pixelStride

        // Bulk-copy ByteBuffers to ByteArrays to avoid per-pixel ByteBuffer.get() overhead.
        val yBuf = yPlane.buffer
        val uBuf = uPlane.buffer
        val vBuf = vPlane.buffer

        val yBytes = ByteArray(yBuf.remaining())
        yBuf.get(yBytes)
        val uBytes = ByteArray(uBuf.remaining())
        uBuf.get(uBytes)
        val vBytes = ByteArray(vBuf.remaining())
        vBuf.get(vBytes)

        val argb = IntArray(w * h)

        for (j in 0 until h) {
            val yRowOffset = j * yRowStride
            val uvRowOffset = (j shr 1) * uvRowStride
            for (i in 0 until w) {
                val y = yBytes[yRowOffset + i].toInt() and 0xFF
                val uvIdx = uvRowOffset + (i shr 1) * uvPixelStride
                val u = (uBytes[uvIdx].toInt() and 0xFF) - 128
                val v = (vBytes[uvIdx].toInt() and 0xFF) - 128

                // Fixed-point: multiply by scaled constant, shift right by 10.
                // 1.402 * 1024 = 1436, 0.344 * 1024 = 352, 0.714 * 1024 = 731, 1.772 * 1024 = 1815
                var r = y + ((v * 1436) shr 10)
                var g = y - ((u * 352 + v * 731) shr 10)
                var b = y + ((u * 1815) shr 10)
                if (r < 0) r = 0 else if (r > 255) r = 255
                if (g < 0) g = 0 else if (g > 255) g = 255
                if (b < 0) b = 0 else if (b > 255) b = 255

                argb[j * w + i] = (0xFF shl 24) or (r shl 16) or (g shl 8) or b
            }
        }

        val bmp = Bitmap.createBitmap(w, h, Bitmap.Config.ARGB_8888)
        bmp.setPixels(argb, 0, w, 0, 0, w, h)
        return bmp
    }
}

// Cached decoder — persists across extractFrame() calls for the same video.
private var cachedDecoder: VideoFrameDecoder? = null
private var cachedDecoderPath: String? = null
private val decoderMutex = Mutex()

actual suspend fun extractFrame(
    videoPath: String, frameTimestamp: Long
): InputFrame? = withContext(Dispatchers.IO) {
    decoderMutex.withLock {
        try {
            if (cachedDecoderPath != videoPath) {
                cachedDecoder?.release()
                cachedDecoder = VideoFrameDecoder(videoPath)
                cachedDecoderPath = videoPath
            }
            cachedDecoder!!.decodeUpTo(frameTimestamp)?.let {
                InputFrame(bitmap = it, timestamp = frameTimestamp)
            }
        } catch (ce: CancellationException) {
            throw ce
        } catch (e: Exception) {
            Logger.w { "extractFrame failed, releasing decoder: ${e.message}" }
            cachedDecoder?.release()
            cachedDecoder = null
            cachedDecoderPath = null
            null
        }
    }
}

/**
 * Batch sequential frame extraction as a Flow. Opens a dedicated decoder,
 * decodes frames in order, and emits each as it's ready. More efficient than
 * repeated [extractFrame] calls because there's no per-call mutex/coroutine
 * overhead and the decoder runs continuously.
 */
actual fun extractFrames(
    videoPath: String, frameTimestamps: List<Long>, cropRect: CropRect?
): Flow<InputFrame> = flow {
    val decoder = VideoFrameDecoder(videoPath)
    try {
        for (ts in frameTimestamps) {
            val bitmap = decoder.decodeUpTo(ts)
            if (bitmap != null) {
                val output = if (cropRect != null) {
                    val cx = cropRect.x.coerceIn(0, bitmap.width - 1)
                    val cy = cropRect.y.coerceIn(0, bitmap.height - 1)
                    val cw = cropRect.width.coerceAtMost(bitmap.width - cx)
                    val ch = cropRect.height.coerceAtMost(bitmap.height - cy)
                    Bitmap.createBitmap(bitmap, cx, cy, cw, ch)
                } else bitmap
                emit(InputFrame(bitmap = output, timestamp = ts))
            }
        }
    } finally {
        decoder.release()
    }
}.flowOn(Dispatchers.IO)

actual suspend fun listVideoFrameTimestamps(
    videoPath: String,
): List<Long> = withContext(Dispatchers.IO) {
    val ctx = VideoExtractionContext.get()
    val uri = videoPath.toUri()
    val extractor = MediaExtractor()

    val out = ArrayList<Long>()
    try {
        extractor.setDataSource(ctx, uri, null)

        var videoTrack = -1
        for (i in 0 until extractor.trackCount) {
            val format = extractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME) ?: continue
            if (mime.startsWith("video/")) {
                videoTrack = i
                break
            }
        }

        if (videoTrack == -1) {
            Logger.i { "No video track found in file: $videoPath" }
            return@withContext out
        }

        extractor.selectTrack(videoTrack)

        while (true) {
            val sampleTimeUs = extractor.sampleTime
            if (sampleTimeUs < 0) break
            out.add(sampleTimeUs / 1000L)
            if (!extractor.advance()) break
        }
    } catch (e: Exception) {
        Logger.i { "Exception extracting frames from $videoPath: ${e::class.simpleName} - ${e.message}" }
    } catch (t: Throwable) {
        Logger.i { "CRASH prevented during frame extraction for $videoPath: ${t::class.simpleName}" }
    } finally {
        extractor.release()
    }
    out
}

actual object VideoExtractionContext {
    private lateinit var application: Application

    fun setUp(context: Context) {
        application = context as Application
    }

    fun get(): Context {
        if (VideoExtractionContext::application.isInitialized.not()) throw Exception("Application context isn't initialized")
        return application.applicationContext
    }
}
