package de.duffy.screencast

import java.io.ByteArrayOutputStream

/**
 * Minimaler MPEG-TS-Muxer fuer den Zweit-Ausgang unter /stream.ts.
 *
 * Damit laesst sich der Stream in VLC, mpv oder Kodi oeffnen - praktisch fuer
 * Fernseher, deren Browser kein MSE koennen, und fuer den Beamer am Laptop.
 */
class TsMuxer(private val sps: ByteArray, private val pps: ByteArray) {

    private val continuity = HashMap<Int, Int>()

    /** Liefert die TS-Pakete fuer einen Frame. Vor Keyframes stehen PAT/PMT. */
    fun frame(nalUnits: List<ByteArray>, ptsUs: Long, keyframe: Boolean): ByteArray {
        val out = ByteArrayOutputStream()
        if (keyframe) {
            out.write(pat())
            out.write(pmt())
        }

        val es = ByteArrayOutputStream()
        // Access Unit Delimiter - manche Player brauchen ihn zur Frame-Abgrenzung.
        es.write(byteArrayOf(0, 0, 0, 1, 9, 0xF0.toByte()))
        if (keyframe) {
            es.write(START_CODE); es.write(sps)
            es.write(START_CODE); es.write(pps)
        }
        for (nal in nalUnits) {
            val type = nal[0].toInt() and 0x1F
            if (type == 7 || type == 8 || type == 9) continue
            es.write(START_CODE); es.write(nal)
        }
        out.write(pes(es.toByteArray(), ptsUs, keyframe))
        return out.toByteArray()
    }

    /** PAT/PMT einzeln - wird an neue HTTP-Clients vorab geschickt. */
    fun tables(): ByteArray {
        val out = ByteArrayOutputStream()
        out.write(pat()); out.write(pmt())
        return out.toByteArray()
    }

    // ------------------------------------------------------------- Tabellen

    private fun pat(): ByteArray {
        val s = ByteArrayOutputStream()
        s.u8(0x00)                    // table_id
        s.u16(0xB00D)                 // syntax indicator + section_length 13
        s.u16(1)                      // transport_stream_id
        s.u8(0xC1)                    // version 0, current
        s.u8(0x00); s.u8(0x00)        // section_number, last_section_number
        s.u16(1)                      // program_number
        s.u16(0xE000 or PMT_PID)      // reserved + PMT PID
        return section(PAT_PID, s.toByteArray())
    }

    private fun pmt(): ByteArray {
        val s = ByteArrayOutputStream()
        s.u8(0x02)
        s.u16(0xB012)                 // section_length 18
        s.u16(1)                      // program_number
        s.u8(0xC1)
        s.u8(0x00); s.u8(0x00)
        s.u16(0xE000 or VIDEO_PID)    // PCR PID
        s.u16(0xF000)                 // program_info_length 0
        s.u8(0x1B)                    // stream_type H.264
        s.u16(0xE000 or VIDEO_PID)
        s.u16(0xF000)                 // ES_info_length 0
        return section(PMT_PID, s.toByteArray())
    }

    /** Packt eine PSI-Sektion samt CRC in genau ein TS-Paket. */
    private fun section(pid: Int, body: ByteArray): ByteArray {
        val withCrc = ByteArrayOutputStream()
        withCrc.write(body)
        withCrc.u32(crc32(body))

        val packet = ByteArray(188)
        packet[0] = 0x47
        packet[1] = (0x40 or (pid shr 8)).toByte()   // payload unit start
        packet[2] = (pid and 0xFF).toByte()
        packet[3] = (0x10 or nextCc(pid)).toByte()   // nur Payload
        packet[4] = 0                                 // pointer_field
        val data = withCrc.toByteArray()
        System.arraycopy(data, 0, packet, 5, data.size)
        for (i in (5 + data.size) until 188) packet[i] = 0xFF.toByte()
        return packet
    }

    // ------------------------------------------------------------------ PES

    private fun pes(es: ByteArray, ptsUs: Long, keyframe: Boolean): ByteArray {
        val pts = ptsUs * 90 / 1000
        val header = ByteArrayOutputStream()
        header.write(byteArrayOf(0, 0, 1, 0xE0.toByte()))
        header.u16(0)                 // Laenge 0 = unbegrenzt, bei Video erlaubt
        header.u8(0x80)               // '10' + keine Flags
        header.u8(0x80)               // nur PTS vorhanden
        header.u8(5)                  // Header-Datenlaenge
        writePts(header, pts)
        val payload = ByteArrayOutputStream()
        payload.write(header.toByteArray())
        payload.write(es)

        val data = payload.toByteArray()
        val out = ByteArrayOutputStream()
        var offset = 0
        var first = true
        while (offset < data.size) {
            val remaining = data.size - offset
            val packet = ByteArray(188)
            packet[0] = 0x47
            packet[1] = (((if (first) 0x40 else 0x00)) or (VIDEO_PID shr 8)).toByte()
            packet[2] = (VIDEO_PID and 0xFF).toByte()

            // Adaptation Field nur wenn noetig: beim ersten Paket (PCR / Random
            // Access) oder am Ende zum Auffuellen.
            var afLength = -1
            if (first) {
                afLength = 7                                  // Flags + PCR
            }
            var payloadSpace = 184 - (if (afLength >= 0) afLength + 1 else 0)
            if (remaining < payloadSpace) {
                // Rest passt: Adaptation Field auf Stuffing erweitern.
                val needed = payloadSpace - remaining
                afLength = if (afLength >= 0) afLength + needed else (needed - 1).coerceAtLeast(0)
                payloadSpace = remaining
            }

            if (afLength >= 0) {
                packet[3] = (0x30 or nextCc(VIDEO_PID)).toByte()  // AF + Payload
                var p = 4
                packet[p++] = afLength.toByte()
                if (afLength > 0) {
                    val flags = (if (keyframe && first) 0x40 else 0x00) or (if (first) 0x10 else 0x00)
                    packet[p++] = flags.toByte()
                    if (first) {
                        writePcr(packet, p, pts); p += 6
                    }
                    while (p < 4 + 1 + afLength) packet[p++] = 0xFF.toByte()
                }
            } else {
                packet[3] = (0x10 or nextCc(VIDEO_PID)).toByte()
            }

            val start = 188 - payloadSpace
            System.arraycopy(data, offset, packet, start, payloadSpace)
            offset += payloadSpace
            first = false
            out.write(packet)
        }
        return out.toByteArray()
    }

    private fun writePts(out: ByteArrayOutputStream, pts: Long) {
        out.u8((0x21 or (((pts shr 30) and 0x07).toInt() shl 1)))
        out.u8((((pts shr 22) and 0xFF).toInt()))
        out.u8((((pts shr 14) and 0x7F).toInt() shl 1) or 1)
        out.u8((((pts shr 7) and 0xFF).toInt()))
        out.u8((((pts and 0x7F).toInt() shl 1) or 1))
    }

    private fun writePcr(packet: ByteArray, offset: Int, pts: Long) {
        val base = pts
        packet[offset] = ((base shr 25) and 0xFF).toByte()
        packet[offset + 1] = ((base shr 17) and 0xFF).toByte()
        packet[offset + 2] = ((base shr 9) and 0xFF).toByte()
        packet[offset + 3] = ((base shr 1) and 0xFF).toByte()
        packet[offset + 4] = ((((base and 1) shl 7) or 0x7E)).toByte()
        packet[offset + 5] = 0
    }

    private fun nextCc(pid: Int): Int {
        val value = (continuity[pid] ?: 0)
        continuity[pid] = (value + 1) and 0x0F
        return value
    }

    companion object {
        const val PAT_PID = 0x0000
        const val PMT_PID = 0x1000
        const val VIDEO_PID = 0x0100
        private val START_CODE = byteArrayOf(0, 0, 0, 1)

        /** CRC-32/MPEG-2 ueber die Sektion. */
        fun crc32(data: ByteArray): Long {
            var crc = 0xFFFFFFFFL
            for (b in data) {
                crc = crc xor ((b.toLong() and 0xFF) shl 24)
                repeat(8) {
                    crc = if (crc and 0x80000000L != 0L) ((crc shl 1) xor 0x04C11DB7L) else (crc shl 1)
                    crc = crc and 0xFFFFFFFFL
                }
            }
            return crc
        }
    }
}
