package de.duffy.screencast

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.os.Build
import android.util.Log
import android.view.Surface
import java.nio.ByteBuffer

/**
 * Hardware-H.264-Encoder. Bekommt das Bildschirmbild ueber eine Surface vom
 * VirtualDisplay und liefert fertige Frames an den [Listener].
 *
 * Die Einstellungen sind auf niedrige Latenz getrimmt: keine B-Frames,
 * konstante Bitrate, kurze Keyframe-Abstaende.
 */
class ScreenEncoder(
    private val width: Int,
    private val height: Int,
    private val bitrate: Int,
    private val frameRate: Int,
    private val listener: Listener
) {
    interface Listener {
        /** Einmalig, sobald der Encoder SPS/PPS ausgegeben hat. */
        fun onConfig(sps: ByteArray, pps: ByteArray, width: Int, height: Int)
        fun onFrame(nalUnits: List<ByteArray>, ptsUs: Long, keyframe: Boolean)
    }

    private var codec: MediaCodec? = null
    private var thread: Thread? = null
    @Volatile private var running = false
    private var firstPtsUs = -1L
    private var lastPtsUs = -1L

    lateinit var surface: Surface
        private set

    fun start() {
        val encoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
        try {
            encoder.configure(buildFormat(true), null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        } catch (e: Exception) {
            // Manche Chips lehnen ein fest vorgegebenes Profil ab - dann ohne.
            Log.w(TAG, "Encoder mag die Profilvorgabe nicht, zweiter Versuch", e)
            encoder.reset()
            encoder.configure(buildFormat(false), null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        }
        surface = encoder.createInputSurface()
        encoder.start()
        codec = encoder

        running = true
        thread = Thread({ drain(encoder) }, "screen-encoder").also { it.start() }
    }

    private fun buildFormat(withProfile: Boolean): MediaFormat {
        return MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height).apply {
            setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface)
            setInteger(MediaFormat.KEY_BIT_RATE, bitrate)
            setInteger(MediaFormat.KEY_FRAME_RATE, frameRate)
            setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 2)
            setInteger(MediaFormat.KEY_BITRATE_MODE, MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_CBR)
            if (withProfile) {
                // Baseline ohne B-Frames: das versteht jeder Fernseher-Browser.
                setInteger(MediaFormat.KEY_PROFILE, MediaCodecInfo.CodecProfileLevel.AVCProfileBaseline)
                setInteger(MediaFormat.KEY_LEVEL, MediaCodecInfo.CodecProfileLevel.AVCLevel4)
            }
            // Standbild trotzdem weitersenden, sonst laeuft der Puffer im Player leer.
            setLong(MediaFormat.KEY_REPEAT_PREVIOUS_FRAME_AFTER, 250_000L)
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
                setInteger(MediaFormat.KEY_PRIORITY, 0) // Echtzeit
            }
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                setInteger(MediaFormat.KEY_LATENCY, 1)
            }
        }
    }

    /** Neuen Keyframe anfordern - noetig, wenn ein Client frisch verbindet. */
    fun requestKeyFrame() {
        try {
            codec?.setParameters(android.os.Bundle().apply {
                putInt(MediaCodec.PARAMETER_KEY_REQUEST_SYNC_FRAME, 0)
            })
        } catch (e: Exception) {
            Log.w(TAG, "Keyframe-Anforderung fehlgeschlagen", e)
        }
    }

    fun stop() {
        running = false
        thread?.join(1500)
        thread = null
        try { codec?.stop() } catch (_: Exception) {}
        try { codec?.release() } catch (_: Exception) {}
        codec = null
        try { if (::surface.isInitialized) surface.release() } catch (_: Exception) {}
    }

    private fun drain(encoder: MediaCodec) {
        val info = MediaCodec.BufferInfo()
        var sps: ByteArray? = null
        var pps: ByteArray? = null

        while (running) {
            val index = try {
                encoder.dequeueOutputBuffer(info, 100_000L)
            } catch (e: IllegalStateException) {
                Log.w(TAG, "Encoder beendet", e); break
            }
            if (index < 0) continue

            val buffer: ByteBuffer? = encoder.getOutputBuffer(index)
            if (buffer == null || info.size <= 0) {
                encoder.releaseOutputBuffer(index, false); continue
            }
            val data = ByteArray(info.size)
            buffer.position(info.offset)
            buffer.get(data, 0, info.size)
            encoder.releaseOutputBuffer(index, false)

            val nals = Fmp4Muxer.splitNalUnits(data, 0, data.size)

            if (info.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG != 0) {
                for (nal in nals) {
                    when (nal[0].toInt() and 0x1F) {
                        7 -> sps = nal
                        8 -> pps = nal
                    }
                }
                if (sps != null && pps != null) listener.onConfig(sps!!, pps!!, width, height)
                continue
            }

            // Manche Encoder liefern SPS/PPS zusammen mit dem ersten Keyframe.
            if (sps == null || pps == null) {
                for (nal in nals) {
                    when (nal[0].toInt() and 0x1F) {
                        7 -> sps = nal
                        8 -> pps = nal
                    }
                }
                if (sps != null && pps != null) listener.onConfig(sps!!, pps!!, width, height)
            }

            if (firstPtsUs < 0) firstPtsUs = info.presentationTimeUs
            var pts = info.presentationTimeUs - firstPtsUs
            if (pts <= lastPtsUs) pts = lastPtsUs + 1000  // monoton halten
            lastPtsUs = pts

            val keyframe = info.flags and MediaCodec.BUFFER_FLAG_KEY_FRAME != 0
            if (nals.isNotEmpty()) listener.onFrame(nals, pts, keyframe)
        }
    }

    companion object {
        private const val TAG = "ScreenEncoder"
    }
}