package com.performancecoachlab.posedetection.encoding

import co.touchlab.kermit.Logger
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.PixelMap
import androidx.compose.ui.graphics.toPixelMap
import kotlinx.cinterop.ByteVar
import kotlinx.cinterop.ExperimentalForeignApi
import kotlinx.cinterop.addressOf
import kotlinx.cinterop.alloc
import kotlinx.cinterop.convert
import kotlinx.cinterop.memScoped
import kotlinx.cinterop.plus
import kotlinx.cinterop.ptr
import kotlinx.cinterop.reinterpret
import kotlinx.cinterop.usePinned
import kotlinx.cinterop.value
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import platform.Foundation.NSThread
import platform.AVFoundation.AVAssetWriter
import platform.AVFoundation.AVAssetWriterInput
import platform.AVFoundation.AVAssetWriterInputPixelBufferAdaptor
import platform.AVFoundation.AVAssetWriterStatusWriting
import platform.AVFoundation.AVFileTypeMPEG4
import platform.AVFoundation.AVMediaTypeVideo
import platform.AVFoundation.AVVideoAverageBitRateKey
import platform.AVFoundation.AVVideoCodecH264
import platform.AVFoundation.AVVideoCodecKey
import platform.AVFoundation.AVVideoCompressionPropertiesKey
import platform.AVFoundation.AVVideoHeightKey
import platform.AVFoundation.AVVideoProfileLevelH264HighAutoLevel
import platform.AVFoundation.AVVideoProfileLevelKey
import platform.AVFoundation.AVVideoWidthKey
import platform.CoreFoundation.kCFAllocatorDefault
import platform.CoreGraphics.CGBitmapContextCreate
import platform.CoreGraphics.CGColorSpaceCreateDeviceRGB
import platform.CoreGraphics.CGColorSpaceRelease
import platform.CoreGraphics.CGContextClearRect
import platform.CoreGraphics.CGContextRelease
import platform.CoreGraphics.CGImageAlphaInfo
import platform.CoreGraphics.CGRectMake
import platform.CoreMedia.CMTimeMake
import platform.CoreVideo.CVPixelBufferGetBaseAddress
import platform.CoreVideo.CVPixelBufferGetBytesPerRow
import platform.CoreVideo.CVPixelBufferLockBaseAddress
import platform.CoreVideo.CVPixelBufferPoolCreatePixelBuffer
import platform.CoreVideo.CVPixelBufferRefVar
import platform.CoreVideo.CVPixelBufferRelease
import platform.CoreVideo.CVPixelBufferUnlockBaseAddress
import platform.CoreVideo.kCVPixelBufferCGBitmapContextCompatibilityKey
import platform.CoreVideo.kCVPixelBufferCGImageCompatibilityKey
import platform.CoreVideo.kCVPixelBufferHeightKey
import platform.CoreVideo.kCVPixelBufferPixelFormatTypeKey
import platform.CoreVideo.kCVPixelBufferWidthKey
import platform.CoreVideo.kCVPixelFormatType_32BGRA
import platform.CoreVideo.kCVReturnSuccess
import platform.Foundation.CFBridgingRelease
import platform.Foundation.NSFileManager
import platform.Foundation.NSURL
import kotlin.native.runtime.NativeRuntimeApi

private class IOSVideoBuilder(
    private val outputPath: String,
    private val fps: Int,
    private val width: Int,
    private val height: Int
) : VideoBuilder {
    private var writer: AVAssetWriter? = null
    private var input: AVAssetWriterInput? = null
    private var adaptor: AVAssetWriterInputPixelBufferAdaptor? = null
    private var frameCount = 0L
    private var started = false

    @OptIn(ExperimentalForeignApi::class)
    private fun initWriter() {
        // Initialize writer and session
        val fileManager = NSFileManager.defaultManager
        if (fileManager.fileExistsAtPath(outputPath)) {
            fileManager.removeItemAtPath(outputPath, null)
        }
        val videoWriter = AVAssetWriter(NSURL.fileURLWithPath(outputPath), AVFileTypeMPEG4, null)
        val outputSettings = mapOf(
            AVVideoCodecKey.toString() to AVVideoCodecH264,
            AVVideoWidthKey.toString() to width,
            AVVideoHeightKey.toString() to height,
            AVVideoCompressionPropertiesKey.toString() to mapOf(
                AVVideoAverageBitRateKey.toString() to 6000000,
                AVVideoProfileLevelKey.toString() to AVVideoProfileLevelH264HighAutoLevel
            )
        )
        val videoWriterInput = AVAssetWriterInput(
            mediaType = AVMediaTypeVideo, outputSettings = outputSettings as Map<Any?, *>?
        )
        val sourcePixelBufferAttributes = mapOf(
            CFBridgingRelease(kCVPixelBufferPixelFormatTypeKey) as String to kCVPixelFormatType_32BGRA,
            CFBridgingRelease(kCVPixelBufferWidthKey) as String to width,
            CFBridgingRelease(kCVPixelBufferHeightKey) as String to height,
            CFBridgingRelease(kCVPixelBufferCGImageCompatibilityKey) as String to true,
            CFBridgingRelease(kCVPixelBufferCGBitmapContextCompatibilityKey) as String to true
        )
        val pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(
            assetWriterInput = videoWriterInput,
            sourcePixelBufferAttributes = sourcePixelBufferAttributes as Map<Any?, *>?
        )
        if (videoWriter.canAddInput(videoWriterInput)) {
            videoWriter.addInput(videoWriterInput)
        } else {
            throw IllegalStateException("Cannot add input to writer")
        }
        videoWriter.startWriting()
        videoWriter.startSessionAtSourceTime(CMTimeMake(0, 1))
        writer = videoWriter
        input = videoWriterInput
        adaptor = pixelBufferAdaptor
        started = true
        frameCount = 0L
    }

    private var lastTimestampMs = -1L

    @OptIn(ExperimentalForeignApi::class, NativeRuntimeApi::class)
    override suspend fun addFrame(frame: ImageBitmap, timestampms: Long) {
        if (!started) initWriter()
        if (frame.width == 0 || frame.height == 0) return
        // Ensure strictly monotonic timestamps.
        val ts = if (timestampms <= lastTimestampMs) lastTimestampMs + 1 else timestampms
        lastTimestampMs = ts
        val videoWriterInput = input!!
        val pixelBufferAdaptor = adaptor!!
        val pool = pixelBufferAdaptor.pixelBufferPool
        if (pool == null) {
            Logger.w { "VideoBuilder: pool null at frame $frameCount" }
            return
        }
        val presentationTime = CMTimeMake(ts, 1_000)
        memScoped {
            val pixelBufferPtr = alloc<CVPixelBufferRefVar>()
            val status =
                CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, pool, pixelBufferPtr.ptr)
            if (status == kCVReturnSuccess) {
                val pixelBuffer = pixelBufferPtr.value
                CVPixelBufferLockBaseAddress(pixelBuffer, 0u)
                val baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer)
                val bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
                val colorSpace = CGColorSpaceCreateDeviceRGB()
                val bitmapInfo = CGImageAlphaInfo.kCGImageAlphaPremultipliedFirst.value
                val context = CGBitmapContextCreate(
                    baseAddress,
                    width.convert(),
                    height.convert(),
                    8u,
                    bytesPerRow.convert(),
                    colorSpace,
                    bitmapInfo
                )
                if (context != null) {
                    CGContextClearRect(
                        context,
                        CGRectMake(0.0, 0.0, width.toDouble(), height.toDouble())
                    )
                    val pixelMap = frame.toPixelMap()
                    val buffer = ensureFourChannelBuffer(pixelMap)
                    val contextData = baseAddress?.reinterpret<ByteVar>()
                    if (contextData != null) {
                        val srcRowBytes = width * 4
                        val dstRowBytes = bytesPerRow.toInt()
                        for (row in 0 until height) {
                            buffer.usePinned { pinned ->
                                val srcOffset = row * srcRowBytes
                                val dstOffset = row * dstRowBytes
                                val dstPtr = contextData.plus(dstOffset)
                                platform.posix.memcpy(
                                    dstPtr,
                                    pinned.addressOf(srcOffset),
                                    srcRowBytes.toULong()
                                )
                            }
                        }
                    }
                    CGContextRelease(context)
                }
                CGColorSpaceRelease(colorSpace)
                CVPixelBufferUnlockBaseAddress(pixelBuffer, 0u)
                // Wait for input to be ready (with timeout guard)
                var waitCount = 0
                while (!videoWriterInput.readyForMoreMediaData) {
                    waitCount++
                    if (waitCount > 10_000_000) {
                        Logger.w { "VideoBuilder: addFrame timeout at frame $frameCount" }
                        CVPixelBufferUnlockBaseAddress(pixelBuffer, 0u)
                        CVPixelBufferRelease(pixelBuffer)
                        return
                    }
                }
                pixelBufferAdaptor.appendPixelBuffer(
                    pixelBuffer,
                    withPresentationTime = presentationTime
                )
                CVPixelBufferRelease(pixelBuffer)
            } else {
                //println("Failed to allocate pixel buffer: $status")
            }
        }
        frameCount++
    }

    @OptIn(ExperimentalForeignApi::class)
    override suspend fun finalize(): String {
        if (!started) return outputPath
        Logger.d { "VideoBuilder: finalize() starting, frameCount=$frameCount" }
        val inp = input
        val w = writer
        if (inp == null || w == null) {
            cleanup()
            return outputPath
        }
        try {
            inp.markAsFinished()
            Logger.d { "VideoBuilder: input marked finished" }

            // finishWritingWithCompletionHandler must be called from main thread.
            // The completion handler fires on AVFoundation's internal queue.
            // Use a truly empty handler and poll w.status instead.
            w.finishWritingWithCompletionHandler {
                // Intentionally empty - do NOT reference any Kotlin objects here
                // as this runs on an AVFoundation internal dispatch queue
            }
            Logger.d { "VideoBuilder: finishWriting called, waiting..." }
            var waitCount = 0
            while (w.status == AVAssetWriterStatusWriting && waitCount < 3000) {
                NSThread.sleepForTimeInterval(0.01)
                waitCount++
            }
            Logger.d { "VideoBuilder: done, status=${w.status}, waited=${waitCount * 10}ms" }
        } catch (e: Exception) {
            Logger.e(e) { "VideoBuilder: finalize exception" }
        }
        if (w.status == platform.AVFoundation.AVAssetWriterStatusFailed) {
            Logger.e { "VideoBuilder: finalize FAILED: ${w.error?.localizedDescription}" }
        } else {
            Logger.d { "VideoBuilder: finalize SUCCESS, frames=$frameCount" }
        }
        cleanup()
        return outputPath
    }

    override suspend fun cancel() = withContext(Dispatchers.Default) {
        writer?.cancelWriting()
        cleanup()
    }

    private fun cleanup() {
        writer = null
        input = null
        adaptor = null
        started = false
        frameCount = 0L
    }

    private fun ensureFourChannelBuffer(pixelMap: PixelMap): ByteArray {
        val width = pixelMap.width
        val height = pixelMap.height
        val buffer = pixelMap.buffer
        // Compose PixelMap's buffer is an IntArray (ARGB for each pixel)
        // If buffer size matches width*height, treat as 4-channel ARGB and convert to BGRA bytes
        if (buffer.size == width * height) {
            val dst = ByteArray(width * height * 4)
            for (i in 0 until width * height) {
                val argb = buffer[i]
                val a = (argb shr 24) and 0xFF
                val r = (argb shr 16) and 0xFF
                val g = (argb shr 8) and 0xFF
                val b = argb and 0xFF
                dst[i * 4] = b.toByte() // B
                dst[i * 4 + 1] = g.toByte() // G
                dst[i * 4 + 2] = r.toByte() // R
                dst[i * 4 + 3] = a.toByte() // A
            }
            return dst
        }
        // If buffer is not ARGB IntArray, fallback (should not happen)
        throw IllegalArgumentException("Unsupported pixel buffer format for ImageBitmap/PixelMap")
    }
}

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

