package de.duffy.screencast.miracast

import de.duffy.screencast.Fmp4Muxer
import de.duffy.screencast.TsMuxer
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertTrue
import org.junit.Test
import java.io.BufferedInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.net.DatagramPacket
import java.net.DatagramSocket
import java.net.InetAddress
import java.net.ServerSocket
import java.net.Socket
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicReference

/**
 * Spielt einen Miracast-Fernseher nach: derselbe RTSP-Ablauf M1 bis M7, den
 * ein echtes Geraet fuehrt. So laesst sich pruefen, ob unsere Seite den
 * Handshake korrekt beherrscht und danach brauchbares Bild schickt.
 */
class MiracastTest {

    @Test
    fun ceaMaskIsParsedAndBestFormatChosen() {
        // So antwortet ein handelsueblicher Fernseher.
        val value = "00 00 02 04 0001DEFF 155557FF 00000FFF 00 0000 0000 11 none none"
        val caps = WfdFormats.parseVideoFormats(value)
        assertEquals(1, caps.size)
        assertEquals(0x02, caps[0].profile)
        assertEquals(0x04, caps[0].level)
        assertEquals(0x0001DEFFL, caps[0].ceaMask)

        val format = WfdFormats.choose(caps)
        assertNotNull(format)
        assertEquals(1280, format!!.width)
        assertEquals(720, format.height)
        assertEquals(30, format.fps)
    }

    @Test
    fun onlyOneFormatBitIsAnnouncedBack() {
        val format = WfdFormats.CEA[5]
        val selection = WfdFormats.buildSelection(format, 0x02, 0x04)
        val fields = selection.split(" ")
        assertEquals(13, fields.size)
        val cea = fields[4].toLong(16)
        assertEquals(1L shl 5, cea)
        assertEquals(java.lang.Long.bitCount(cea), 1)
        assertEquals("none", fields[11])
    }

    @Test
    fun fallsBackWhenSinkOnlySupportsSmallFormats() {
        // Nur 640x480p60 (Bit 0) und 720x480p60 (Bit 1).
        val caps = WfdFormats.parseVideoFormats("00 00 01 02 00000003 00000000 00000000 00 0000 0000 00 none none")
        val format = WfdFormats.choose(caps)
        assertEquals(640, format!!.width)
        assertEquals(480, format.height)
    }

    @Test
    fun rejectsWhenNoCommonFormat() {
        val caps = WfdFormats.parseVideoFormats("00 00 01 02 00000000 00000000 00000000 00 0000 0000 00 none none")
        assertEquals(null, WfdFormats.choose(caps))
    }

    /**
     * Der grosse Durchlauf: kompletter Handshake gegen den simulierten
     * Fernseher, danach echtes Videomaterial als RTP-Strom.
     */
    @Test
    fun fullHandshakeAndStream() {
        val server = ServerSocket(0)
        val sinkPort = server.localPort
        val transcript = java.util.Collections.synchronizedList(ArrayList<String>())
        val playing = CountDownLatch(1)
        val chosenFormat = AtomicReference<WfdFormats.VideoFormat>()
        val rtpPortRef = AtomicReference<Int>()

        // ---- Der simulierte Fernseher -------------------------------------
        val udp = DatagramSocket()
        val sinkThread = Thread {
            server.accept().use { client ->
                val input = BufferedInputStream(client.getInputStream())
                val out = client.getOutputStream()

                // Ein echter Fernseher liest laufend mit und sortiert selbst,
                // was Anfrage und was Antwort ist - Reihenfolge nicht garantiert.
                val requests = java.util.concurrent.LinkedBlockingQueue<RtspMessage>()
                val responses = java.util.concurrent.LinkedBlockingQueue<RtspMessage>()
                Thread {
                    while (true) {
                        val msg = RtspMessage.read(input) ?: break
                        if (msg.isRequest) requests.put(msg) else responses.put(msg)
                    }
                }.apply { isDaemon = true }.start()
                fun nextRequest(): RtspMessage =
                    requests.poll(10, TimeUnit.SECONDS) ?: error("Keine Anfrage erhalten")
                fun nextResponse(): RtspMessage =
                    responses.poll(10, TimeUnit.SECONDS) ?: error("Keine Antwort erhalten")

                // M1: Wir erwarten OPTIONS von der Quelle.
                val m1 = nextRequest()
                transcript.add("M1 ${m1.method}")
                assertEquals("OPTIONS", m1.method)
                assertEquals("org.wfa.wfd1.0", m1.headers["require"])
                out.write(RtspMessage.response(m1.cseq, mapOf(
                    "Public" to "org.wfa.wfd1.0, SETUP, TEARDOWN, PLAY, PAUSE, GET_PARAMETER, SET_PARAMETER"
                )))
                out.flush()

                // M2: Der Fernseher fragt zurueck.
                out.write(RtspMessage.request("OPTIONS", "*", 100,
                    mapOf("Require" to "org.wfa.wfd1.0")))
                out.flush()
                val m2 = nextResponse()
                transcript.add("M2 Antwort ${m2.statusCode}")
                assertEquals(200, m2.statusCode)
                assertTrue(m2.headers["public"]!!.contains("org.wfa.wfd1.0"))

                // M3: Faehigkeiten werden abgefragt.
                val m3 = nextRequest()
                transcript.add("M3 ${m3.method}")
                assertEquals("GET_PARAMETER", m3.method)
                assertTrue(m3.body.contains("wfd_video_formats"))
                out.write(RtspMessage.response(m3.cseq, body = buildString {
                    append("wfd_video_formats: 00 00 02 04 0001DEFF 155557FF 00000FFF 00 0000 0000 11 none none\r\n")
                    append("wfd_audio_codecs: AAC 00000007 00\r\n")
                    append("wfd_client_rtp_ports: RTP/AVP/UDP;unicast ${udp.localPort} 0 mode=play\r\n")
                    append("wfd_content_protection: none\r\n")
                }))
                out.flush()

                // M4: Die Quelle legt das Format fest.
                val m4 = nextRequest()
                transcript.add("M4 ${m4.method}")
                assertEquals("SET_PARAMETER", m4.method)
                val announced = m4.param("wfd_video_formats")!!
                val ceaBits = announced.split(" ")[4].toLong(16)
                assertEquals("Genau ein Format muss angekündigt werden",
                    1, java.lang.Long.bitCount(ceaBits))
                assertTrue(m4.body.contains("wfd_presentation_URL"))
                out.write(RtspMessage.response(m4.cseq)); out.flush()

                // M5: Aufforderung, die Sitzung einzurichten.
                val m5 = nextRequest()
                transcript.add("M5 ${m5.param("wfd_trigger_method")}")
                assertEquals("SETUP", m5.param("wfd_trigger_method"))
                out.write(RtspMessage.response(m5.cseq)); out.flush()

                // M6: Der Fernseher richtet die Sitzung ein und nennt seinen Port.
                out.write(RtspMessage.request("SETUP",
                    "rtsp://127.0.0.1/wfd1.0/streamid=0", 101,
                    mapOf("Transport" to "RTP/AVP/UDP;unicast;client_port=${udp.localPort}")))
                out.flush()
                val m6 = nextResponse()
                transcript.add("M6 Antwort ${m6.statusCode}")
                assertEquals(200, m6.statusCode)
                assertNotNull("Session-Kopfzeile fehlt", m6.headers["session"])
                assertTrue("server_port fehlt", m6.headers["transport"]!!.contains("server_port="))

                // M7: Wiedergabe anfordern.
                out.write(RtspMessage.request("PLAY",
                    "rtsp://127.0.0.1/wfd1.0/streamid=0", 102,
                    mapOf("Session" to m6.session!!)))
                out.flush()
                val m7 = nextResponse()
                transcript.add("M7 Antwort ${m7.statusCode}")
                assertEquals(200, m7.statusCode)
                Thread.sleep(4000)
            }
        }
        sinkThread.isDaemon = true
        sinkThread.start()

        // ---- Unsere Seite --------------------------------------------------
        val socket = Socket("127.0.0.1", sinkPort)
        val session = RtspSession(socket, "192.168.49.1", 1080, object : RtspSession.Listener {
            override fun onFormatChosen(format: WfdFormats.VideoFormat) { chosenFormat.set(format) }
            override fun onPlay(rtpPort: Int, format: WfdFormats.VideoFormat) {
                rtpPortRef.set(rtpPort); playing.countDown()
            }
            override fun onStopped(reason: String) { transcript.add("Ende: $reason") }
            override fun onLog(message: String) { }
        })
        session.localRtpPort = 55555
        session.start()

        assertTrue("Handshake kam nicht bis PLAY: $transcript", playing.await(15, TimeUnit.SECONDS))
        assertEquals(1280, chosenFormat.get().width)
        assertEquals(udp.localPort, rtpPortRef.get())
        println("Handshake-Ablauf: ${transcript.joinToString(" | ")}")

        // ---- Bild schicken und beim Fernseher wieder einsammeln ------------
        val sample = File(System.getProperty("sampleH264")!!).readBytes()
        val nals = Fmp4Muxer.splitNalUnits(sample, 0, sample.size)
        var sps: ByteArray? = null; var pps: ByteArray? = null
        val frames = 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)) {
                        frames.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) frames.add(Pair(current, key))

        val received = ByteArrayOutputStream()
        var gaps = 0
        var packets = 0
        val collector = Thread {
            val buffer = ByteArray(2048)
            // Bis der Sender fertig ist: erst grosszuegig warten, dann
            // abbrechen, sobald eine halbe Sekunde nichts mehr kommt. Nur so
            // liegt am Ende ein vollstaendiger Strom vor - ein mittendrin
            // abgeschnittener Frame liesse sich sonst nicht dekodieren.
            udp.soTimeout = 5000
            udp.receiveBufferSize = 4 * 1024 * 1024
            var expected = -1
            try {
                while (true) {
                    val packet = DatagramPacket(buffer, buffer.size)
                    udp.receive(packet)
                    packets++
                    // RTP-Kopf pruefen und Nutzlast herausloesen
                    assertEquals(0x80.toByte(), packet.data[0])
                    assertEquals(33.toByte(), packet.data[1])
                    assertTrue("Nutzlast muss ein Vielfaches von 188 sein",
                        (packet.length - 12) % 188 == 0)
                    assertEquals("TS-Paket muss mit 0x47 beginnen", 0x47.toByte(), packet.data[12])
                    val sequence = ((packet.data[2].toInt() and 0xFF) shl 8) or (packet.data[3].toInt() and 0xFF)
                    if (expected >= 0 && sequence != expected) gaps++
                    expected = (sequence + 1) and 0xFFFF
                    received.write(packet.data, 12, packet.length - 12)
                    udp.soTimeout = 500
                }
            } catch (_: Exception) { /* Zeitueberschreitung = Strom zu Ende */ }
        }
        collector.start()

        val muxer = TsMuxer(sps!!, pps!!)
        val sender = RtpTsSender(InetAddress.getByName("127.0.0.1"), udp.localPort)
        var pts = 0L
        for ((frameNals, isKey) in frames) {
            sender.send(muxer.frame(frameNals, pts, isKey), pts)
            pts += 33_333L
            Thread.sleep(3)
        }
        collector.join(15000)
        sender.close()
        session.stop()

        val out = File(System.getProperty("muxOutDir") ?: "build/muxtest").apply { mkdirs() }
            .resolve("miracast.ts")
        out.writeBytes(received.toByteArray())
        println("Über RTP empfangen: ${out.length()} Bytes in $packets Paketen, $gaps Sequenzlücken")
        println("Datei: ${out.absolutePath}")
        assertEquals("RTP-Sequenznummern müssen lückenlos sein", 0, gaps)
        assertTrue("Zu wenig Bilddaten angekommen", out.length() > 30_000)
        assertEquals("Der Strom muss aus ganzen TS-Paketen bestehen", 0L, out.length() % 188L)
    }
}

/**
 * Prueft die Richtung, die der Wi-Fi-Display-Standard vorsieht: Die Quelle
 * hoert zu, der Fernseher ruft an. Genau daran ist die erste Fassung
 * gescheitert - der Fernseher wies unseren Anruf ab (ECONNREFUSED).
 */
class MiracastDirectionTest {

    @org.junit.Test
    fun sourceListensAndStartsHandshakeWhenTvCallsIn() {
        val port = 17236
        val log = java.util.Collections.synchronizedList(ArrayList<String>())
        val pipeline = MiracastPipeline(
            projection = null,
            densityDpi = 320,
            handler = android.os.Handler(),
            maxHeight = 1080,
            bitrate = 6_000_000,
            listener = object : MiracastPipeline.Listener {
                override fun onStatus(text: String) { log.add(text) }
                override fun onStreaming(format: WfdFormats.VideoFormat) { log.add("streaming $format") }
                override fun onFinished(reason: String) { log.add("ende: $reason") }
            }
        )
        try {
            pipeline.startListening(listOf(port))
            Thread.sleep(300)

            // Ab hier spielen wir den Fernseher, der sich bei uns meldet.
            java.net.Socket("127.0.0.1", port).use { tv ->
                val input = java.io.BufferedInputStream(tv.getInputStream())
                val out = tv.getOutputStream()

                val m1 = RtspMessage.read(input)
                assertNotNull("Die Quelle muss von sich aus M1 schicken", m1)
                assertEquals("OPTIONS", m1!!.method)
                assertEquals("org.wfa.wfd1.0", m1.headers["require"])

                out.write(RtspMessage.response(m1.cseq, mapOf(
                    "Public" to "org.wfa.wfd1.0, SETUP, PLAY, GET_PARAMETER, SET_PARAMETER"
                )))
                out.flush()

                // Als naechstes muessen die Faehigkeiten abgefragt werden.
                tv.soTimeout = 8000
                var m3: RtspMessage? = null
                while (m3 == null) {
                    val msg = RtspMessage.read(input) ?: break
                    if (msg.isRequest && msg.method == "GET_PARAMETER") m3 = msg
                    else if (msg.isRequest) out.write(RtspMessage.response(msg.cseq)).also { out.flush() }
                }
                assertNotNull("Nach M1 muss M3 folgen", m3)
                assertTrue(m3!!.body.contains("wfd_video_formats"))
            }
            assertTrue("Es muss auf Port $port gelauscht werden: $log",
                log.any { it.contains("Warte auf den Fernseher") })
            assertTrue("Die eingehende Verbindung muss gemeldet werden: $log",
                log.any { it.contains("hat sich verbunden") })
        } finally {
            pipeline.stop()
        }
    }
}
