package com.performancecoachlab.posedetection.recording

import android.app.Application
import android.content.Context
import android.media.MediaExtractor
import android.media.MediaFormat
import android.media.MediaMetadataRetriever
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import androidx.core.net.toUri
import co.touchlab.kermit.Logger

actual suspend fun extractFrame(
    videoPath: String, frameTimestamp: Long
): InputFrame? {
    return withContext(Dispatchers.IO) {
        val retriever = MediaMetadataRetriever()
        try {
            // Attempt to set the data source and handle any exceptions gracefully
            val uri = videoPath.toUri()

            retriever.setDataSource(VideoExtractionContext.get(),uri)
            try {
                val bitmap = retriever.getFrameAtTime(
                    frameTimestamp * 1000L, // microseconds
                    MediaMetadataRetriever.OPTION_CLOSEST
                )
                bitmap?.also {
                    return@withContext InputFrame(bitmap = it, timestamp = frameTimestamp)
                }
            } catch (e: Exception) {
                e.printStackTrace()
            }
        } catch (e: Exception) {
            // Handle any exceptions during the setup phase (e.g., invalid URL)
            e.printStackTrace()
        } finally {
            // Ensure that the retriever is released regardless of success or failure
            retriever.release()
        }
        return@withContext null
    }
}

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)

        // Select video track
        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)

        var idx = 0
        while (true) {
            val sampleTimeUs = extractor.sampleTime
            if (sampleTimeUs < 0) break

            out.add(sampleTimeUs / 1000L) // ms
            if (out.size >= Int.MAX_VALUE) break
            idx++

            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
    }
}