package de.duffy.screencast

import android.util.Log
import java.util.concurrent.CopyOnWriteArrayList

/**
 * Verteilt die encodierten Frames an alle verbundenen Zuschauer.
 *
 * Der Encoder laeuft genau einmal, egal wie viele Geraete zuschauen. Jeder
 * Client bekommt beim Verbinden das Init-Segment und steigt beim naechsten
 * Keyframe ein.
 */
class StreamHub {

    /** Ein Zuschauer - entweder Browser (WebSocket) oder Player (MPEG-TS). */
    interface Client {
        /** Puffert Daten fuer diesen Client. false = Client haengt und wurde geleert. */
        fun enqueue(data: ByteArray): Boolean
        fun close()
    }

    interface MediaClient : Client {
        fun sendStreamInfo(codec: String, width: Int, height: Int)
        fun sendReset()
    }

    private val browserClients = CopyOnWriteArrayList<MediaClient>()
    private val tsClients = CopyOnWriteArrayList<Client>()

    @Volatile private var mp4: Fmp4Muxer? = null
    @Volatile private var ts: TsMuxer? = null
    @Volatile private var initSegment: ByteArray? = null
    @Volatile private var lastFramePtsUs = 0L

    /** Wird gesetzt, damit ein neu verbundener Client sofort ein Bild bekommt. */
    var onKeyFrameNeeded: (() -> Unit)? = null

    val viewerCount: Int get() = browserClients.size + tsClients.size

    // Warten diese Clients noch auf ihren ersten Keyframe?
    private val waitingForKey = HashSet<Client>()

    @Synchronized
    fun onEncoderConfig(sps: ByteArray, pps: ByteArray, width: Int, height: Int) {
        val muxer = Fmp4Muxer(sps, pps, width, height)
        mp4 = muxer
        ts = TsMuxer(sps, pps)
        initSegment = muxer.initSegment()
        // Laufende Browser-Clients neu aufsetzen (z.B. nach Drehung des Handys).
        for (client in browserClients) {
            client.sendReset()
            client.sendStreamInfo(muxer.codecString, width, height)
            client.enqueue(initSegment!!)
            synchronized(waitingForKey) { waitingForKey.add(client) }
        }
        onKeyFrameNeeded?.invoke()
    }

    @Synchronized
    fun onFrame(nalUnits: List<ByteArray>, ptsUs: Long, keyframe: Boolean) {
        val duration = (ptsUs - lastFramePtsUs).coerceIn(1000L, 500_000L)
        lastFramePtsUs = ptsUs

        if (keyframe) {
            synchronized(waitingForKey) { waitingForKey.clear() }
        }

        if (browserClients.isNotEmpty()) {
            mp4?.let { muxer ->
                val fragment = muxer.fragment(nalUnits, ptsUs, duration, keyframe)
                if (fragment.isNotEmpty()) {
                    for (client in browserClients) {
                        if (isWaiting(client)) continue
                        if (!client.enqueue(fragment)) {
                            // Rueckstand wurde verworfen: erst beim naechsten
                            // Keyframe wieder einsteigen, sonst ruckelt es.
                            synchronized(waitingForKey) { waitingForKey.add(client) }
                            onKeyFrameNeeded?.invoke()
                        }
                    }
                }
            }
        }

        if (tsClients.isNotEmpty()) {
            ts?.let { muxer ->
                val packets = muxer.frame(nalUnits, ptsUs, keyframe)
                for (client in tsClients) {
                    if (isWaiting(client)) continue
                    if (!client.enqueue(packets)) {
                        synchronized(waitingForKey) { waitingForKey.add(client) }
                        onKeyFrameNeeded?.invoke()
                    }
                }
            }
        }
    }

    private fun isWaiting(client: Client): Boolean =
        synchronized(waitingForKey) { waitingForKey.contains(client) }

    fun addBrowserClient(client: MediaClient) {
        val muxer = mp4
        val init = initSegment
        if (muxer != null && init != null) {
            client.sendStreamInfo(muxer.codecString, muxer.width, muxer.height)
            client.enqueue(init)
        }
        synchronized(waitingForKey) { waitingForKey.add(client) }
        browserClients.add(client)
        onKeyFrameNeeded?.invoke()
        Log.i(TAG, "Browser verbunden, jetzt $viewerCount Zuschauer")
    }

    fun addTsClient(client: Client): Boolean {
        val muxer = ts ?: return false
        client.enqueue(muxer.tables())
        synchronized(waitingForKey) { waitingForKey.add(client) }
        tsClients.add(client)
        onKeyFrameNeeded?.invoke()
        return true
    }

    fun removeClient(client: Client) {
        browserClients.remove(client)
        tsClients.remove(client)
        synchronized(waitingForKey) { waitingForKey.remove(client) }
        Log.i(TAG, "Client getrennt, jetzt $viewerCount Zuschauer")
    }

    fun closeAll() {
        for (c in browserClients) c.close()
        for (c in tsClients) c.close()
        browserClients.clear(); tsClients.clear()
        synchronized(waitingForKey) { waitingForKey.clear() }
    }

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