/*
 * Copyright 2023 LiveKit, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package io.livekit.android.room

import android.javax.sdp.MediaDescription
import android.javax.sdp.SdpFactory
import dagger.assisted.Assisted
import dagger.assisted.AssistedFactory
import dagger.assisted.AssistedInject
import io.livekit.android.dagger.InjectionNames
import io.livekit.android.room.util.*
import io.livekit.android.util.Either
import io.livekit.android.util.LKLog
import io.livekit.android.util.debounce
import io.livekit.android.webrtc.SdpExt
import io.livekit.android.webrtc.SdpFmtp
import io.livekit.android.webrtc.getExts
import io.livekit.android.webrtc.getFmtps
import io.livekit.android.webrtc.getMsid
import io.livekit.android.webrtc.getRtps
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import org.webrtc.*
import org.webrtc.PeerConnection.RTCConfiguration
import javax.inject.Named
import kotlin.math.roundToLong

/**
 * @suppress
 */
internal class PeerConnectionTransport
@AssistedInject
constructor(
    @Assisted config: RTCConfiguration,
    @Assisted pcObserver: PeerConnection.Observer,
    @Assisted private val listener: Listener?,
    @Named(InjectionNames.DISPATCHER_IO)
    private val ioDispatcher: CoroutineDispatcher,
    connectionFactory: PeerConnectionFactory,
    private val sdpFactory: SdpFactory,
) {
    private val coroutineScope = CoroutineScope(ioDispatcher + SupervisorJob())
    internal val peerConnection: PeerConnection = connectionFactory.createPeerConnection(
        config,
        pcObserver,
    ) ?: throw IllegalStateException("peer connection creation failed?")
    private val pendingCandidates = mutableListOf<IceCandidate>()
    private var restartingIce: Boolean = false

    private var renegotiate = false

    private val mutex = Mutex()

    private var trackBitrates = mutableMapOf<TrackBitrateInfoKey, TrackBitrateInfo>()

    interface Listener {
        fun onOffer(sd: SessionDescription)
    }

    fun addIceCandidate(candidate: IceCandidate) {
        runBlocking {
            mutex.withLock {
                if (peerConnection.remoteDescription != null && !restartingIce) {
                    peerConnection.addIceCandidate(candidate)
                } else {
                    pendingCandidates.add(candidate)
                }
            }
        }
    }

    suspend fun setRemoteDescription(sd: SessionDescription): Either<Unit, String?> {
        val result = peerConnection.setRemoteDescription(sd)
        if (result is Either.Left) {
            mutex.withLock {
                pendingCandidates.forEach { pending ->
                    peerConnection.addIceCandidate(pending)
                }
                pendingCandidates.clear()
                restartingIce = false
            }
        }

        if (this.renegotiate) {
            this.renegotiate = false
            this.createAndSendOffer()
        }

        return result
    }

    val negotiate = debounce<MediaConstraints?, Unit>(100, coroutineScope) {
        if (it != null) {
            createAndSendOffer(it)
        } else {
            createAndSendOffer()
        }
    }

    suspend fun createAndSendOffer(constraints: MediaConstraints = MediaConstraints()) {
        if (listener == null) {
            return
        }

        val iceRestart =
            constraints.findConstraint(MediaConstraintKeys.ICE_RESTART) == MediaConstraintKeys.TRUE
        if (iceRestart) {
            LKLog.d { "restarting ice" }
            restartingIce = true
        }

        if (this.peerConnection.signalingState() == PeerConnection.SignalingState.HAVE_LOCAL_OFFER) {
            // we're waiting for the peer to accept our offer, so we'll just wait
            // the only exception to this is when ICE restart is needed
            val curSd = peerConnection.remoteDescription
            if (iceRestart && curSd != null) {
                // TODO: handle when ICE restart is needed but we don't have a remote description
                // the best thing to do is to recreate the peerconnection
                peerConnection.setRemoteDescription(curSd)
            } else {
                renegotiate = true
                return
            }
        }

        // actually negotiate
        LKLog.d { "starting to negotiate" }
        val sdpOffer = when (val outcome = peerConnection.createOffer(constraints)) {
            is Either.Left -> outcome.value
            is Either.Right -> {
                LKLog.d { "error creating offer: ${outcome.value}" }
                return
            }
        }

        // munge sdp
        val sdpDescription = sdpFactory.createSessionDescription(sdpOffer.description)

        val mediaDescs = sdpDescription.getMediaDescriptions(true)
        for (mediaDesc in mediaDescs) {
            if (mediaDesc !is MediaDescription) {
                continue
            }
            if (mediaDesc.media.mediaType == "audio") {
                // TODO
            } else if (mediaDesc.media.mediaType == "video") {
                ensureVideoDDExtensionForSVC(mediaDesc)
                ensureCodecBitrates(mediaDesc, trackBitrates = trackBitrates)
            }
        }

        val finalSdp = setMungedSdp(sdpOffer, sdpDescription.toString())
        listener.onOffer(finalSdp)
    }

    private suspend fun setMungedSdp(sdp: SessionDescription, mungedDescription: String, remote: Boolean = false): SessionDescription {
        val mungedSdp = SessionDescription(sdp.type, mungedDescription)

        LKLog.v { "sdp type: ${sdp.type}\ndescription:\n${sdp.description}" }
        LKLog.v { "munged sdp type: ${mungedSdp.type}\ndescription:\n${mungedSdp.description}" }
        val mungedResult = if (remote) {
            peerConnection.setRemoteDescription(mungedSdp)
        } else {
            peerConnection.setLocalDescription(mungedSdp)
        }

        val mungedErrorMessage = when (mungedResult) {
            is Either.Left -> {
                // munged sdp set successfully.
                return mungedSdp
            }

            is Either.Right -> {
                if (mungedResult.value.isNullOrBlank()) {
                    "unknown sdp error"
                } else {
                    mungedResult.value
                }
            }
        }

        // munged sdp setting failed
        LKLog.w {
            "setting munged sdp for " +
                "${if (remote) "remote" else "local"} description, " +
                "${mungedSdp.type} type failed, falling back to unmodified."
        }
        LKLog.w { "error: $mungedErrorMessage" }

        val result = if (remote) {
            peerConnection.setRemoteDescription(sdp)
        } else {
            peerConnection.setLocalDescription(sdp)
        }

        if (result is Either.Right) {
            val errorMessage = if (result.value.isNullOrBlank()) {
                "unknown sdp error"
            } else {
                result.value
            }

            // sdp setting failed
            LKLog.w {
                "setting original sdp for " +
                    "${if (remote) "remote" else "local"} description, " +
                    "${sdp.type} type failed!"
            }
            LKLog.w { "error: $errorMessage" }
        }
        return sdp
    }

    fun prepareForIceRestart() {
        restartingIce = true
    }

    fun close() {
        peerConnection.dispose()
    }

    fun updateRTCConfig(config: RTCConfiguration) {
        peerConnection.setConfiguration(config)
    }

    fun registerTrackBitrateInfo(cid: String, trackBitrateInfo: TrackBitrateInfo) {
        trackBitrates[TrackBitrateInfoKey.Cid(cid)] = trackBitrateInfo
    }

    fun registerTrackBitrateInfo(transceiver: RtpTransceiver, trackBitrateInfo: TrackBitrateInfo) {
        trackBitrates[TrackBitrateInfoKey.Transceiver(transceiver)] = trackBitrateInfo
    }

    @AssistedFactory
    interface Factory {
        fun create(
            config: RTCConfiguration,
            pcObserver: PeerConnection.Observer,
            listener: Listener?,
        ): PeerConnectionTransport
    }
}

private const val DD_EXTENSION_URI = "https://aomediacodec.github.io/av1-rtp-spec/#dependency-descriptor-rtp-header-extension"

internal fun ensureVideoDDExtensionForSVC(mediaDesc: MediaDescription) {
    val codec = mediaDesc.getRtps()
        .firstOrNull()
        ?.second
        ?.codec ?: return
    if (!isSVCCodec(codec)) {
        return
    }

    var maxId = 0L

    val ddFound = mediaDesc.getExts().any { (_, ext) ->
        if (ext.uri == DD_EXTENSION_URI) {
            return@any true
        }
        if (ext.value > maxId) {
            maxId = ext.value
        }
        false
    }

    // Not found, add manually
    if (!ddFound) {
        mediaDesc.addAttribute(
            SdpExt(
                value = maxId + 1,
                uri = DD_EXTENSION_URI,
                config = null,
                direction = null,
                encryptUri = null,
            ).toAttributeField(),
        )
    }
}

/* The svc codec (av1/vp9) would use a very low bitrate at the begining and
increase slowly by the bandwidth estimator until it reach the target bitrate. The
process commonly cost more than 10 seconds cause subscriber will get blur video at
the first few seconds. So we use a 70% of target bitrate here as the start bitrate to
eliminate this issue.
*/
private const val startBitrateForSVC = 0.7

internal fun ensureCodecBitrates(
    media: MediaDescription,
    trackBitrates: Map<TrackBitrateInfoKey, TrackBitrateInfo>,
) {
    val msid = media.getMsid()?.value ?: return
    for ((key, trackBr) in trackBitrates) {
        if (key !is TrackBitrateInfoKey.Cid) {
            continue
        }

        val (cid) = key
        if (!msid.contains(cid)) {
            continue
        }

        val (_, rtp) = media.getRtps()
            .firstOrNull { (_, rtp) -> rtp.codec.equals(trackBr.codec, ignoreCase = true) }
            ?: continue
        val codecPayload = rtp.payload

        val fmtps = media.getFmtps()
        var fmtpFound = false
        for ((attribute, fmtp) in fmtps) {
            if (fmtp.payload == codecPayload) {
                fmtpFound = true
                var newFmtpConfig = fmtp.config
                if (!fmtp.config.contains("x-google-start-bitrate")) {
                    newFmtpConfig = "$newFmtpConfig;x-google-start-bitrate=${(trackBr.maxBitrate * startBitrateForSVC).roundToLong()}"
                }
                if (!fmtp.config.contains("x-google-max-bitrate")) {
                    newFmtpConfig = "$newFmtpConfig;x-google-max-bitrate=${trackBr.maxBitrate}"
                }
                if (fmtp.config != newFmtpConfig) {
                    attribute.value = "${fmtp.payload} $newFmtpConfig"
                    break
                }
            }
        }

        if (!fmtpFound) {
            media.addAttribute(
                SdpFmtp(
                    payload = codecPayload,
                    config = "x-google-start-bitrate=${trackBr.maxBitrate * startBitrateForSVC};" +
                        "x-google-max-bitrate=${trackBr.maxBitrate}",
                ).toAttributeField(),
            )
        }
    }
}

internal fun isSVCCodec(codec: String?): Boolean {
    return codec != null &&
        ("av1".equals(codec, ignoreCase = true) ||
            "vp9".equals(codec, ignoreCase = true))
}

internal data class TrackBitrateInfo(
    val codec: String,
    val maxBitrate: Long,
)

sealed class TrackBitrateInfoKey {
    data class Cid(val value: String) : TrackBitrateInfoKey()
    data class Transceiver(val value: RtpTransceiver) : TrackBitrateInfoKey()
}