package com.dronex.rec.recording

import android.media.MediaCodec
import android.media.MediaExtractor
import android.media.MediaFormat
import android.media.MediaMuxer
import java.io.File
import java.io.FileInputStream
import java.io.FileOutputStream
import java.nio.ByteBuffer

object RecordingSegmentMuxer {
    fun mergeOrCopy(
            segments: List<File>,
            outputFile: File,
    ) {
        val validSegments = segments.filter { it.exists() && it.length() > 0L }
        require(validSegments.isNotEmpty()) { "No recording segments were generated" }
        if (validSegments.size == 1) {
            copyFile(validSegments.first(), outputFile)
            return
        }

        MediaMuxer(outputFile.absolutePath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4).use {
                muxer ->
            val firstTracks = readTrackFormats(validSegments.first())
            val videoTrack = firstTracks.video?.let { muxer.addTrack(it) }
            val audioTrack = firstTracks.audio?.let { muxer.addTrack(it) }
            require(videoTrack != null) { "No video track in recording segment" }

            muxer.start()
            var offsetUs = 0L
            validSegments.forEach { segment ->
                val durationUs =
                        copySegmentTracks(
                                segment = segment,
                                muxer = muxer,
                                videoOutputTrack = videoTrack,
                                audioOutputTrack = audioTrack,
                                offsetUs = offsetUs,
                        )
                offsetUs += durationUs + 1L
            }
        }
    }

    private fun copyFile(source: File, target: File) {
        FileInputStream(source).use { input ->
            FileOutputStream(target).use { output -> input.copyTo(output) }
        }
    }

    private fun readTrackFormats(file: File): TrackFormats {
        val extractor = MediaExtractor()
        extractor.setDataSource(file.absolutePath)
        try {
            var video: MediaFormat? = null
            var audio: MediaFormat? = null
            for (index in 0 until extractor.trackCount) {
                val format = extractor.getTrackFormat(index)
                val mime = format.getString(MediaFormat.KEY_MIME).orEmpty()
                when {
                    mime.startsWith("video/") && video == null -> video = format
                    mime.startsWith("audio/") && audio == null -> audio = format
                }
            }
            return TrackFormats(video = video, audio = audio)
        } finally {
            extractor.release()
        }
    }

    private fun copySegmentTracks(
            segment: File,
            muxer: MediaMuxer,
            videoOutputTrack: Int,
            audioOutputTrack: Int?,
            offsetUs: Long,
    ): Long {
        var segmentDurationUs = 0L
        segmentDurationUs =
                maxOf(
                        segmentDurationUs,
                        copyTrack(segment, muxer, "video/", videoOutputTrack, offsetUs),
                )
        if (audioOutputTrack != null) {
            segmentDurationUs =
                    maxOf(
                            segmentDurationUs,
                            copyTrack(segment, muxer, "audio/", audioOutputTrack, offsetUs),
                    )
        }
        return segmentDurationUs
    }

    private fun copyTrack(
            segment: File,
            muxer: MediaMuxer,
            mimePrefix: String,
            outputTrack: Int,
            offsetUs: Long,
    ): Long {
        val extractor = MediaExtractor()
        extractor.setDataSource(segment.absolutePath)
        try {
            val inputTrack = findTrack(extractor, mimePrefix) ?: return 0L
            extractor.selectTrack(inputTrack)
            val bufferSize = trackBufferSize(extractor.getTrackFormat(inputTrack))
            val buffer = ByteBuffer.allocate(bufferSize)
            val info = MediaCodec.BufferInfo()
            var firstSampleTimeUs: Long? = null
            var lastSampleTimeUs = 0L

            while (true) {
                buffer.clear()
                val sampleSize = extractor.readSampleData(buffer, 0)
                if (sampleSize < 0) break

                val sampleTimeUs = extractor.sampleTime
                val baseTimeUs = firstSampleTimeUs ?: sampleTimeUs.also { firstSampleTimeUs = it }
                val normalizedTimeUs = (sampleTimeUs - baseTimeUs).coerceAtLeast(0L)
                info.set(
                        0,
                        sampleSize,
                        offsetUs + normalizedTimeUs,
                        extractor.sampleFlags,
                )
                muxer.writeSampleData(outputTrack, buffer, info)
                lastSampleTimeUs = normalizedTimeUs
                extractor.advance()
            }
            return lastSampleTimeUs
        } finally {
            extractor.release()
        }
    }

    private fun findTrack(extractor: MediaExtractor, mimePrefix: String): Int? {
        for (index in 0 until extractor.trackCount) {
            val mime = extractor.getTrackFormat(index).getString(MediaFormat.KEY_MIME).orEmpty()
            if (mime.startsWith(mimePrefix)) {
                return index
            }
        }
        return null
    }

    private fun trackBufferSize(format: MediaFormat): Int {
        return if (format.containsKey(MediaFormat.KEY_MAX_INPUT_SIZE)) {
            format.getInteger(MediaFormat.KEY_MAX_INPUT_SIZE).coerceAtLeast(DEFAULT_BUFFER_SIZE)
        } else {
            DEFAULT_BUFFER_SIZE
        }
    }

    private data class TrackFormats(
            val video: MediaFormat?,
            val audio: MediaFormat?,
    )

    private const val DEFAULT_BUFFER_SIZE = 2 * 1024 * 1024
}

private inline fun MediaMuxer.use(block: (MediaMuxer) -> Unit) {
    try {
        block(this)
    } finally {
        try {
            stop()
        } catch (_: Exception) {
        }
        release()
    }
}
