package de.duffy.screencast

import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import java.io.BufferedInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.io.InputStream
import java.net.Socket
import java.security.MessageDigest

/**
 * Faehrt Server und Verteiler wirklich hoch und holt den Stream ueber echte
 * Sockets ab - genau wie es der Fernseher spaeter tut. Der WebSocket-Client
 * ist bewusst von Hand geschrieben, damit auch Handshake und Rahmenformat
 * geprueft werden.
 */
class ServerTest {

    private fun frames(): Triple<ByteArray, ByteArray, List<Pair<List<ByteArray>, Boolean>>> {
        val data = File(System.getProperty("sampleH264")!!).readBytes()
        val nals = Fmp4Muxer.splitNalUnits(data, 0, data.size)
        var sps: ByteArray? = null; var pps: ByteArray? = null
        val result = ArrayList<Pair<List<ByteArray>, Boolean>>()
        var current = ArrayList<ByteArray>(); var key = false; var hasSlice = false
        for (nal in nals) {
            when (val type = nal[0].toInt() and 0x1F) {
                7 -> sps = nal
                8 -> pps = nal
                else -> {
                    if (hasSlice && (type == 9 || type == 6 || type == 1 || type == 5)) {
                        result.add(Pair(current, key)); current = ArrayList(); hasSlice = false; key = false
                    }
                    if (type == 1 || type == 5) { hasSlice = true; key = type == 5 }
                    current.add(nal)
                }
            }
        }
        if (hasSlice) result.add(Pair(current, key))
        return Triple(sps!!, pps!!, result)
    }

    private fun outDir(): File = File(System.getProperty("muxOutDir") ?: "build/muxtest").apply { mkdirs() }

    private fun readHeaders(input: InputStream): Pair<String, Map<String, String>> {
        val line = StringBuilder(); val lines = ArrayList<String>()
        while (true) {
            val b = input.read()
            if (b < 0) break
            if (b == '\n'.code) {
                val text = line.toString().trimEnd('\r')
                if (text.isEmpty()) break
                lines.add(text); line.setLength(0)
            } else line.append(b.toChar())
        }
        val headers = HashMap<String, String>()
        for (i in 1 until lines.size) {
            val idx = lines[i].indexOf(':')
            if (idx > 0) headers[lines[i].substring(0, idx).lowercase()] = lines[i].substring(idx + 1).trim()
        }
        return Pair(lines.firstOrNull() ?: "", headers)
    }

    @Test
    fun playerPageIsServed() {
        val hub = StreamHub()
        val server = WebServer({ "<html>hallo tv</html>".toByteArray() }, hub, 18100)
        server.start()
        try {
            Socket("127.0.0.1", server.port).use { socket ->
                socket.getOutputStream().write("GET / HTTP/1.1\r\nHost: x\r\n\r\n".toByteArray())
                val input = BufferedInputStream(socket.getInputStream())
                val (status, headers) = readHeaders(input)
                assertTrue("Status: $status", status.startsWith("HTTP/1.1 200"))
                assertEquals("text/html; charset=utf-8", headers["content-type"])
                val body = ByteArray(headers["content-length"]!!.toInt())
                var read = 0
                while (read < body.size) {
                    val n = input.read(body, read, body.size - read)
                    if (n < 0) break
                    read += n
                }
                assertTrue(String(body).contains("hallo tv"))
            }
        } finally { server.stop() }
    }

    @Test
    fun browserReceivesPlayableFragmentedMp4() {
        val (sps, pps, frames) = frames()
        val hub = StreamHub()
        val server = WebServer({ ByteArray(0) }, hub, 18200)
        server.start()
        try {
            hub.onEncoderConfig(sps, pps, 640, 480)

            Socket("127.0.0.1", server.port).use { socket ->
                val key = "dGhlIHNhbXBsZSBub25jZQ=="
                socket.getOutputStream().write(
                    ("GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n" +
                        "Sec-WebSocket-Key: $key\r\nSec-WebSocket-Version: 13\r\n\r\n").toByteArray()
                )
                val input = BufferedInputStream(socket.getInputStream())
                val (status, headers) = readHeaders(input)
                assertTrue("Status: $status", status.startsWith("HTTP/1.1 101"))

                // Der Accept-Header muss exakt nach RFC 6455 berechnet sein,
                // sonst bricht jeder Browser die Verbindung ab.
                val expected = WebServer.base64(
                    MessageDigest.getInstance("SHA-1")
                        .digest((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").toByteArray())
                )
                assertEquals(expected, headers["sec-websocket-accept"])
                assertEquals("s3pPLMBiTxaQ9kYGzzhZRbK+xOo=", expected)

                // Auf die init-Nachricht warten, damit der Client sicher
                // registriert ist, bevor Frames hereinkommen.
                socket.soTimeout = 5000
                val (firstOpcode, firstPayload) = readWsFrame(input)!!
                assertEquals(0x1, firstOpcode)
                val initText = String(firstPayload)
                assertTrue("init fehlt Codec: $initText", initText.contains("\"codec\":\"avc1."))
                assertTrue(initText.contains("\"width\":640"))
                assertTrue(initText.contains("\"height\":480"))

                var fed = 0
                val feeder = Thread {
                    try {
                        var pts = 0L
                        for ((nals, isKey) in frames) {
                            hub.onFrame(nals, pts, isKey); fed++; pts += 33_333L; Thread.sleep(3)
                        }
                    } catch (e: Throwable) {
                        println("FEEDER-FEHLER bei Frame $fed: $e")
                        e.printStackTrace()
                    }
                    println("Feeder fertig: $fed von ${frames.size} Frames eingespeist")
                }
                feeder.start()

                val media = ByteArrayOutputStream()
                var fragments = 0
                socket.soTimeout = 4000
                try {
                    while (fragments < frames.size) {
                        val (opcode, payload) = readWsFrame(input) ?: break
                        if (opcode == 0x2) { media.write(payload); fragments++ }
                    }
                } catch (e: Exception) {
                    println("Lesen beendet: $e")
                }
                feeder.join(5000)
                println("Empfangen: $fragments Fragmente (+Init), ${media.size()} Bytes von $fed Frames")
                assertTrue("Zu wenig Fragmente: $fragments von ${frames.size}", fragments >= frames.size)

                val out = File(outDir(), "websocket.mp4")
                out.writeBytes(media.toByteArray())
                println("WebSocket lieferte ${out.length()} Bytes -> ${out.absolutePath}")
                assertTrue(out.length() > 25_000)
            }
        } finally { server.stop(); hub.closeAll() }
    }

    /** Liest einen unmaskierten Server-Frame. */
    private fun readWsFrame(input: InputStream): Pair<Int, ByteArray>? {
        val b0 = input.read(); if (b0 < 0) return null
        val opcode = b0 and 0x0F
        val b1 = input.read(); if (b1 < 0) return null
        assertTrue("Server darf nicht maskieren", b1 and 0x80 == 0)
        var length = (b1 and 0x7F).toLong()
        if (length == 126L) length = ((input.read() shl 8) or input.read()).toLong()
        else if (length == 127L) { length = 0; repeat(8) { length = (length shl 8) or input.read().toLong() } }
        val payload = ByteArray(length.toInt())
        var read = 0
        while (read < payload.size) {
            val n = input.read(payload, read, payload.size - read)
            if (n < 0) return null
            read += n
        }
        return Pair(opcode, payload)
    }

    @Test
    fun vlcReceivesPlayableTransportStream() {
        val (sps, pps, frames) = frames()
        val hub = StreamHub()
        val server = WebServer({ ByteArray(0) }, hub, 18300)
        server.start()
        try {
            hub.onEncoderConfig(sps, pps, 640, 480)
            Socket("127.0.0.1", server.port).use { socket ->
                socket.soTimeout = 8000
                socket.getOutputStream().write("GET /stream.ts HTTP/1.1\r\nHost: x\r\n\r\n".toByteArray())
                val input = BufferedInputStream(socket.getInputStream())
                val (status, headers) = readHeaders(input)
                assertTrue("Status: $status", status.startsWith("HTTP/1.1 200"))
                assertEquals("video/mp2t", headers["content-type"])
                Thread.sleep(300) // Client sicher registriert

                val feeder = Thread {
                    var pts = 0L
                    for ((nals, isKey) in frames) {
                        hub.onFrame(nals, pts, isKey); pts += 33_333L; Thread.sleep(3)
                    }
                }
                feeder.start()

                val collected = ByteArrayOutputStream()
                val buffer = ByteArray(8192)
                while (collected.size() < 25_000) {
                    val n = input.read(buffer)
                    if (n < 0) break
                    collected.write(buffer, 0, n)
                }
                feeder.join(5000)
                val out = File(outDir(), "http.ts")
                out.writeBytes(collected.toByteArray())
                println("MPEG-TS lieferte ${out.length()} Bytes -> ${out.absolutePath}")
                assertTrue(collected.toByteArray()[0] == 0x47.toByte())
                assertTrue(out.length() > 25_000)
            }
        } finally { server.stop(); hub.closeAll() }
    }
}
