package com.performancecoachlab.posedetection.encoding

import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.PixelMap
import androidx.compose.ui.graphics.toPixelMap
import com.performancecoachlab.posedetection.recording.InputFrame
import kotlinx.cinterop.ByteVar
import kotlinx.cinterop.CValue
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.useContents
import kotlinx.cinterop.usePinned
import kotlinx.cinterop.value
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
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.CGContextDrawImage
import platform.CoreGraphics.CGContextRelease
import platform.CoreGraphics.CGContextScaleCTM
import platform.CoreGraphics.CGContextTranslateCTM
import platform.CoreGraphics.CGImageAlphaInfo
import platform.CoreGraphics.CGImageRef
import platform.CoreGraphics.CGRectMake
import platform.CoreGraphics.CGSize
import platform.CoreGraphics.CGSizeMake
import platform.CoreGraphics.kCGBitmapByteOrder32Little
import platform.CoreMedia.CMTimeAdd
import platform.CoreMedia.CMTimeMake
import platform.CoreVideo.CVPixelBufferGetBaseAddress
import platform.CoreVideo.CVPixelBufferGetBytesPerRow
import platform.CoreVideo.CVPixelBufferLockBaseAddress
import platform.CoreVideo.CVPixelBufferPoolCreatePixelBuffer
import platform.CoreVideo.CVPixelBufferPoolRef
import platform.CoreVideo.CVPixelBufferRef
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.NSDocumentDirectory
import platform.Foundation.NSFileManager
import platform.Foundation.NSURL
import platform.Foundation.NSUserDomainMask
import platform.darwin.dispatch_queue_create
import kotlin.native.runtime.GC
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
    }
    @OptIn(ExperimentalForeignApi::class, NativeRuntimeApi::class)
    override suspend fun addFrame(frame: ImageBitmap, timestampms: Long) {
        if (!started) initWriter()
        // Check for empty or invalid frame (skip if width or height is 0)
        if (frame.width == 0 || frame.height == 0) return
        val videoWriterInput = input!!
        val pixelBufferAdaptor = adaptor!!
        val pool = pixelBufferAdaptor.pixelBufferPool ?: throw IllegalStateException("Pixel buffer pool is null")
        val presentationTime = CMTimeMake(timestampms, 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
                while (!videoWriterInput.readyForMoreMediaData) {}
                pixelBufferAdaptor.appendPixelBuffer(pixelBuffer, withPresentationTime = presentationTime)
                CVPixelBufferRelease(pixelBuffer)
            } else {
                //println("Failed to allocate pixel buffer: $status")
            }
        }
        frameCount++
    }

    @OptIn(ExperimentalForeignApi::class)
    override suspend fun finalize(): String = withContext(Dispatchers.Default) {
        if (started) {
            input?.markAsFinished()
            writer?.finishWritingWithCompletionHandler {
                //println("[VideoBuilder] Finished writing video!")
            }
            // Wait for writing to finish
            while (writer?.status == AVAssetWriterStatusWriting) {}
            cleanup()
        }
        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)

