from typing import Dict, Literal, List, Optional, Any

import logging
from aiortc import (
    RTCIceServer,
    RTCPeerConnection,
    RTCSessionDescription,
    RTCRtpTransceiver,
    MediaStreamTrack,
    RTCConfiguration,
)
import sdp_transform
from .sdp import common_utils
from .sdp.remote_sdp import RemoteSdp
from .sdp.unified_plan_utils import addLegacySimulcast, getRtpEncodings
from .sdp.common_utils import applyCodecParameters, extractDtlsParameters
from .handler_interface import HandlerInterface
from ..ortc import ExtendedRtpCapabilities
from ..rtp_parameters import (
    MediaKind,
    RtpParameters,
    RtpCapabilities,
    RtpCodecCapability,
    RtpEncodingParameters,
    RtcpParameters,
)
from ..sctp_parameters import SctpCapabilities, SctpParameters, SctpStreamParameters
from ..ortc import getSendingRtpParameters, getSendingRemoteRtpParameters, reduceCodecs
from ..scalability_modes import parse as smParse
from ..models.transport import IceCandidate, IceParameters, DtlsParameters, DtlsRole
from ..models.handler_interface import (
    HandlerRunOptions,
    HandlerSendOptions,
    HandlerSendResult,
    HandlerSendDataChannelResult,
    HandlerReceiveDataChannelResult,
    HandlerReceiveOptions,
    HandlerReceiveResult,
    HandlerReceiveDataChannelOptions,
)
from ..producer import ProducerCodecOptions


logger = logging.getLogger(__name__)
SCTP_NUM_STREAMS = {"OS": 1024, "MIS": 1024}


class AiortcHandler(HandlerInterface):

    def __init__(self, tracks: List[MediaStreamTrack] = [], loop=None):
        super(AiortcHandler, self).__init__(loop=loop)
        # Handler direction.
        self._direction: Optional[Literal["send", "recv"]] = None
        # Remote SDP handler.
        self._remoteSdp: Optional[RemoteSdp] = None
        # Generic sending RTP parameters for audio and video.
        self._sendingRtpParametersByKind: Dict[str, RtpParameters] = {}
        # Generic sending RTP parameters for audio and video suitable for the SDP
        # remote answer.
        self._sendingRemoteRtpParametersByKind: Dict[str, RtpParameters] = {}
        # RTCPeerConnection instance.
        self._pc: Optional[RTCPeerConnection] = None
        # Map of RTCTransceivers indexed by MID.
        self._mapMidTransceiver: Dict[str, RTCRtpTransceiver] = {}
        # Whether a DataChannel m=application section has been created.
        self._hasDataChannelMediaSection = False
        # Sending DataChannel id value counter. Incremented for each new DataChannel.
        self._nextSendSctpStreamId = 0
        # Got transport local and remote parameters.
        self._transportReady = False
        self._tracks = tracks

    @classmethod
    def createFactory(cls, tracks: List[MediaStreamTrack] = [], loop=None):
        return lambda: cls(tracks, loop)

    @property
    def name(self) -> str:
        return "aiortc"

    @property
    def pc(self) -> RTCPeerConnection:
        if self._pc:
            return self._pc
        else:
            raise Exception("PeerConnection not ready")

    @property
    def remoteSdp(self) -> RemoteSdp:
        if self._remoteSdp:
            return self._remoteSdp
        else:
            raise Exception("Remote SDP not ready")

    async def close(self):
        logger.debug("close()")

        if self._pc:
            await self._pc.close()

    async def getNativeRtpCapabilities(self) -> RtpCapabilities:
        logger.debug("getNativeRtpCapabilities()")

        pc = RTCPeerConnection()
        for track in self._tracks:
            pc.addTrack(track)
        pc.addTransceiver("audio")
        pc.addTransceiver("video")

        offer: RTCSessionDescription = await pc.createOffer()
        await pc.close()

        sdpDict: dict = sdp_transform.parse(offer.sdp)
        nativeRtpCapabilities: RtpCapabilities = common_utils.extractRtpCapabilities(
            sdpDict
        )

        return nativeRtpCapabilities

    async def getNativeSctpCapabilities(self) -> SctpCapabilities:
        logger.debug("getNativeSctpCapabilities()")
        return SctpCapabilities.parse_obj({"numStreams": SCTP_NUM_STREAMS})

    def run(
        self,
        direction: Literal["send", "recv"],
        iceParameters: IceParameters,
        iceCandidates: List[IceCandidate],
        dtlsParameters: DtlsParameters,
        extendedRtpCapabilities: ExtendedRtpCapabilities,
        sctpParameters: Optional[SctpParameters] = None,
        iceServers: Optional[RTCIceServer] = None,
        iceTransportPolicy: Optional[Literal["all", "relay"]] = None,
        additionalSettings: Optional[Any] = None,
        proprietaryConstraints: Optional[Any] = None,
    ):
        logger.debug("AiortcHandler run()")
        options = HandlerRunOptions(
            direction=direction,
            iceParameters=iceParameters,
            iceCandidates=iceCandidates,
            dtlsParameters=dtlsParameters,
            sctpParameters=sctpParameters,
            iceServers=iceServers,
            iceTransportPolicy=iceTransportPolicy,
            additionalSettings=additionalSettings,
            proprietaryConstraints=proprietaryConstraints,
            extendedRtpCapabilities=extendedRtpCapabilities,
        )
        self._direction = options.direction
        self._remoteSdp = RemoteSdp(
            iceParameters=options.iceParameters,
            iceCandidates=options.iceCandidates,
            dtlsParameters=options.dtlsParameters,
            sctpParameters=options.sctpParameters,
        )
        self._sendingRtpParametersByKind = {
            "audio": getSendingRtpParameters("audio", options.extendedRtpCapabilities),
            "video": getSendingRtpParameters("video", options.extendedRtpCapabilities),
        }
        self._sendingRemoteRtpParametersByKind = {
            "audio": getSendingRemoteRtpParameters(
                "audio", options.extendedRtpCapabilities
            ),
            "video": getSendingRemoteRtpParameters(
                "video", options.extendedRtpCapabilities
            ),
        }

        configuration = RTCConfiguration(iceServers=iceServers)

        self._pc = RTCPeerConnection(configuration=configuration)

        @self._pc.on("iceconnectionstatechange")
        def on_iceconnectionstatechange():
            if self._pc.iceConnectionState == "checking":
                self.emit("@connectionstatechange", "connecting")
            elif self._pc.iceConnectionState in ["connected", "completed"]:
                self.emit("@connectionstatechange", "connected")
            elif self._pc.iceConnectionState == "failed":
                self.emit("@connectionstatechange", "failed")
            elif self._pc.iceConnectionState == "disconnected":
                self.emit("@connectionstatechange", "disconnected")
            elif self._pc.iceConnectionState == "closed":
                self.emit("@connectionstatechange", "closed")

    async def updateIceServers(self, iceServers):
        logger.warning("updateIceServers() not implemented")
        # TODO: aiortc can not update iceServers

    async def restartIce(self, iceParameters):
        logger.debug("restartIce()")
        self._remoteSdp.updateIceParameters(iceParameters)
        if not self._transportReady:
            return
        if self._direction == "send":
            # NOTE: aiortc RTCPeerConnection createOffer do not have iceRestart options
            offer = await self._pc.createOffer()
            logger.debug(
                f"restartIce() | calling pc.setLocalDescription() [offer:{offer}]"
            )
            await self._pc.setLocalDescription(offer)
            answer: RTCSessionDescription = RTCSessionDescription(
                type="answer", sdp=self._remoteSdp.getSdp()
            )
            logger.debug(
                f"restartIce() | calling pc.setRemoteDescription() [answer:{answer}]"
            )
            await self._pc.setRemoteDescription(answer)
        else:
            offer: RTCSessionDescription = RTCSessionDescription(
                type="offer", sdp=self._remoteSdp.getSdp()
            )
            logger.debug(
                f"restartIce() | calling pc.setRemoteDescription() [offer:{offer}]"
            )
            await self._pc.setRemoteDescription(offer)
            answer = await self._pc.createAnswer()
            logger.debug(
                f"restartIce() | calling pc.setLocalDescription() [answer:{answer}]"
            )
            await self._pc.setLocalDescription(answer)

    async def getTransportStats(self):
        return self._pc.getStats()

    async def send(
        self,
        track: MediaStreamTrack,
        encodings: List[RtpEncodingParameters] = [],
        codecOptions: Optional[ProducerCodecOptions] = None,
        codec: Optional[RtpCodecCapability] = None,
    ) -> HandlerSendResult:
        options = HandlerSendOptions(
            track=track, encodings=encodings, codecOptions=codecOptions, codec=codec
        )
        self._assertSendDirection()
        logger.debug(f"send() [kind:{options.track.kind}, track.id:{options.track.id}]")
        if options.encodings:
            for idx in range(len(options.encodings)):
                options.encodings[idx].rid = f"r{idx}"

        sendingRtpParameters: RtpParameters = self._sendingRtpParametersByKind[
            options.track.kind
        ].copy(deep=True)
        sendingRtpParameters.codecs = reduceCodecs(
            sendingRtpParameters.codecs, options.codec
        )

        sendingRemoteRtpParameters: RtpParameters = (
            self._sendingRemoteRtpParametersByKind[options.track.kind].copy(deep=True)
        )
        sendingRemoteRtpParameters.codecs = reduceCodecs(
            sendingRemoteRtpParameters.codecs, options.codec
        )

        mediaSectionIdx = self.remoteSdp.getNextMediaSectionIdx()
        transceiver = self.pc.addTransceiver(options.track, direction="sendonly")

        offer: RTCSessionDescription = await self.pc.createOffer()
        offerMediaDict: dict
        localSdpDict = sdp_transform.parse(offer.sdp)
        if not self._transportReady:
            await self._setupTransport(
                localDtlsRole="server", localSdpDict=localSdpDict
            )
        # Special case for VP9 with SVC.
        hackVp9Svc = False
        if options.encodings:
            layers = smParse(
                options.encodings[0].scalabilityMode
                if options.encodings[0].scalabilityMode
                else ""
            )
        else:
            layers = smParse("")
        if (
            len(options.encodings) == 1
            and layers.spatialLayers > 1
            and sendingRtpParameters.codecs[0].mimeType.lower() == "video/vp9"
        ):
            logger.debug("send() | enabling legacy simulcast for VP9 SVC")
            hackVp9Svc = True
            localSdpDict = sdp_transform.parse
            offerMediaDict = localSdpDict["media"][mediaSectionIdx.idx]
            addLegacySimulcast(
                offerMediaDict=offerMediaDict, numStreams=layers.spatialLayers
            )
            offer = RTCSessionDescription(
                type="offer", sdp=sdp_transform.write(localSdpDict)
            )

        logger.debug(f"send() | calling pc.setLocalDescription() [offer:{offer}]")

        await self.pc.setLocalDescription(offer)
        # We can now get the transceiver.mid.
        localId = transceiver.mid
        # Set MID.
        sendingRtpParameters.mid = localId
        localSdpDict = sdp_transform.parse(self.pc.localDescription.sdp)

        offerMediaDict = localSdpDict["media"][mediaSectionIdx.idx]

        logger.debug(
            f"send() | get offerMediaDict {offerMediaDict} \n from localSdpDict {localSdpDict['media']} index {mediaSectionIdx.idx}"  # noqa
        )
        # Set RTCP CNAME.
        if sendingRtpParameters.rtcp is None:
            sendingRtpParameters.rtcp = RtcpParameters()
        sendingRtpParameters.rtcp.cname = common_utils.getCname(offerMediaDict)
        # Set RTP encodings by parsing the SDP offer if no encodings are given.
        if not options.encodings:
            sendingRtpParameters.encodings = getRtpEncodings(offerMediaDict)
        # Set RTP encodings by parsing the SDP offer and complete them with given
        # one if just a single encoding has been given.
        elif len(options.encodings) == 1:
            newEncodings = getRtpEncodings(offerMediaDict)
            if newEncodings and options.encodings[0]:
                firstEncodingDict: dict = newEncodings[0].dict()
                optionsEncodingDict: dict = options.encodings[0].dict()
                firstEncodingDict.update(optionsEncodingDict)
                newEncodings[0] = RtpEncodingParameters(**firstEncodingDict)
                if hackVp9Svc:
                    newEncodings = [newEncodings[0]]
            sendingRtpParameters.encodings = newEncodings
        # Otherwise if more than 1 encoding are given use them verbatim.
        else:
            sendingRtpParameters.encodings = options.encodings
        # If VP8 or H264 and there is effective simulcast, add scalabilityMode to
        # each encoding.
        if len(sendingRtpParameters.encodings) > 1 and (
            sendingRtpParameters.codecs[0].mimeType.lower() == "video/vp8"
            or sendingRtpParameters.codecs[0].mimeType.lower() == "video/h264"
        ):
            for encoding in sendingRtpParameters.encodings:
                encoding.scalabilityMode = "S1T3"
        self.remoteSdp.send(
            offerMediaDict=offerMediaDict,
            reuseMid=mediaSectionIdx.reuseMid,
            offerRtpParameters=sendingRtpParameters,
            answerRtpParameters=sendingRemoteRtpParameters,
            codecOptions=options.codecOptions,
            extmapAllowMixed=True,
        )
        answer: RTCSessionDescription = RTCSessionDescription(
            type="answer", sdp=self.remoteSdp.getSdp()
        )
        logger.debug(f"send() | calling pc.setRemoteDescription() [answer:{answer}]")
        await self.pc.setRemoteDescription(answer)
        # Store in the map.
        self._mapMidTransceiver[localId] = transceiver
        return HandlerSendResult(
            localId=localId,
            rtpParameters=sendingRtpParameters,
            rtpSender=transceiver.sender,
        )

    async def stopSending(self, localId):
        pass
        # self._assertSendDirection()
        # logger.debug(f'stopSending() [localId:{localId}]')
        # transceiver = self._mapMidTransceiver.get(localId)
        # if not transceiver:
        #     raise Exception('associated RTCRtpTransceiver not found')
        # transceiver.sender.replaceTrack()
        # # TODO:RTCPeerConnection do not have removeTrack()
        # self._pc.removeTrack(transceiver.sender)
        # self._remoteSdp.closeMediaSection(transceiver.mid)
        # offer = await self._pc.createOffer()
        # logger.debug(f'stopSending() | calling pc.setLocalDescription() [offer:{offer}]')
        # await self._pc.localDescription(offer)
        # answer: RTCSessionDescription = RTCSessionDescription(
        #     type='answer',
        #     sdp=self._remoteSdp.getSdp()
        # )
        # logger.debug(f'stopSending() | calling pc.setRemoteDescription() [answer:{answer}]')
        # await self._pc.setRemoteDescription(answer)
        # self._mapMidTransceiver.pop(localId, None)

    async def replaceTrack(self, localId, track=None):
        self._assertSendDirection()
        if track:
            logger.debug(f"replaceTrack() [localId:{localId}, track.id:{track.id}]")
        else:
            logger.debug(f"replaceTrack() [localId:{localId}, no track]")
        transceiver = self._mapMidTransceiver.get(localId)
        if not transceiver:
            raise Exception("associated RTCRtpTransceiver not found")

        await transceiver.sender.replaceTrack(track)

    async def setMaxSpatialLayer(self, localId: str, spatialLayer: int):
        logger.warning("setMaxSpatialLayer() not implemented")
        # NOTE: RTCRtpSender do not have getParameters()
        # self._assertSendDirection()
        # logger.debug(f'setMaxSpatialLayer() [localId:{localId}, spatialLayer:{spatialLayer}]')
        # transceiver = self._mapMidTransceiver.get(localId)
        # if not transceiver:
        #     raise Exception('associated RTCRtpTransceiver not found')
        # parameters = transceiver.sender.getParameters()

    async def setRtpEncodingParameters(self, localId: str, params: Any):
        logger.warning("setRtpEncodingParameters() not implemented")
        # NOTE: RTCRtpSender do not have getParameters()

    async def getSenderStats(self, localId: str):
        self._assertSendDirection()
        transceiver = self._mapMidTransceiver.get(localId)
        if not transceiver:
            raise Exception("associated RTCRtpTransceiver not found")
        return await transceiver.sender.getStats()

    async def sendDataChannel(
        self,
        streamId: Optional[int] = None,
        ordered: Optional[bool] = True,
        maxPacketLifeTime: Optional[int] = None,
        maxRetransmits: Optional[int] = None,
        label: Optional[str] = None,
        protocol: Optional[str] = None,
    ) -> HandlerSendDataChannelResult:
        if streamId is None:
            streamId = self._nextSendSctpStreamId
        options = SctpStreamParameters(
            streamId=streamId,
            ordered=ordered,
            maxPacketLifeTime=maxPacketLifeTime,
            maxRetransmits=maxRetransmits,
            label=label,
            protocol=protocol,
        )
        self._assertSendDirection()
        logger.debug("sendDataChannel()")
        dataChannel = self.pc.createDataChannel(
            label=options.label,
            maxPacketLifeTime=options.maxPacketLifeTime,
            ordered=options.ordered,
            protocol=options.protocol,
            negotiated=True,
            id=self._nextSendSctpStreamId,
        )
        # Increase next id.
        self._nextSendSctpStreamId = (
            self._nextSendSctpStreamId + 1
        ) % SCTP_NUM_STREAMS.get("MIS", 1)
        # If this is the first DataChannel we need to create the SDP answer with
        # m=application section.
        if not self._hasDataChannelMediaSection:
            offer: RTCSessionDescription = await self.pc.createOffer()
            localSdpDict = sdp_transform.parse(offer.sdp)
            offerMediaDicts = [
                m for m in localSdpDict.get("media") if m.get("type") == "application"
            ]
            if not offerMediaDicts:
                raise Exception("No datachannel")
            offerMediaDict = offerMediaDicts[0]

            if not self._transportReady:
                await self._setupTransport(
                    localDtlsRole="server", localSdpDict=localSdpDict
                )

            logger.debug(
                f"sendDataChannel() | calling pc.setLocalDescription() [offer:{offer}]"
            )
            await self.pc.setLocalDescription(offer)
            self.remoteSdp.sendSctpAssociation(offerMediaDict=offerMediaDict)
            answer: RTCSessionDescription = RTCSessionDescription(
                type="answer", sdp=self.remoteSdp.getSdp()
            )

            logger.debug(
                f"sendDataChannel() | calling pc.setRemoteDescription() [answer:{answer}]"
            )
            await self.pc.setRemoteDescription(answer)
            self._hasDataChannelMediaSection = True

        return HandlerSendDataChannelResult(
            dataChannel=dataChannel, sctpStreamParameters=options
        )

    async def receive(
        self, trackId: str, kind: MediaKind, rtpParameters: RtpParameters
    ) -> HandlerReceiveResult:
        options = HandlerReceiveOptions(
            trackId=trackId, kind=kind, rtpParameters=rtpParameters
        )
        self._assertRecvDirection()
        logger.debug(f"receive() [trackId:{options.trackId}, kind:{options.kind}]")
        localId = (
            options.rtpParameters.mid
            if options.rtpParameters.mid is not None
            else str(len(self._mapMidTransceiver))
        )
        self.remoteSdp.receive(
            mid=localId,
            kind=options.kind,
            offerRtpParameters=options.rtpParameters,
            streamId=options.rtpParameters.rtcp.cname,
            trackId=options.trackId,
        )
        offer: RTCSessionDescription = RTCSessionDescription(
            type="offer", sdp=self.remoteSdp.getSdp()
        )
        logger.debug(f"receive() | calling pc.setRemoteDescription() [offer:{offer}]")
        await self.pc.setRemoteDescription(offer)
        answer: RTCSessionDescription = await self.pc.createAnswer()
        localSdpDict = sdp_transform.parse(answer.sdp)
        answerMediaDict = [
            m for m in localSdpDict.get("media") if str(m.get("mid")) == localId
        ][0]
        # May need to modify codec parameters in the answer based on codec
        # parameters in the offer.
        applyCodecParameters(
            offerRtpParameters=options.rtpParameters, answerMediaDict=answerMediaDict
        )
        answer = RTCSessionDescription(
            type="answer", sdp=sdp_transform.write(localSdpDict)
        )
        if not self._transportReady:
            await self._setupTransport(
                localDtlsRole="client", localSdpDict=localSdpDict
            )
        logger.debug(f"receive() | calling pc.setLocalDescription() [answer:{answer}]")
        await self.pc.setLocalDescription(answer)
        transceivers = [t for t in self.pc.getTransceivers() if t.mid == localId]
        if not transceivers:
            raise Exception("new RTCRtpTransceiver not found")
        # Store in the map.
        transceiver = transceivers[0]
        self._mapMidTransceiver[localId] = transceiver

        return HandlerReceiveResult(
            localId=localId,
            track=transceiver.receiver.track,
            rtpReceiver=transceiver.receiver,
        )

    async def stopReceiving(self, localId: str):
        self._assertRecvDirection()
        logger.debug(f"stopReceiving() [localId:{localId}]")
        transceiver = self._mapMidTransceiver.get(localId)
        if not transceiver:
            raise Exception("associated RTCRtpTransceiver not found")
        self.remoteSdp.closeMediaSection(transceiver.mid)
        offer: RTCSessionDescription = RTCSessionDescription(
            type="offer", sdp=self.remoteSdp.getSdp()
        )
        logger.debug(
            f"stopReceiving() | calling pc.setRemoteDescription() [offer:{offer}]"
        )
        await self.pc.setRemoteDescription(offer)
        answer = await self.pc.createAnswer()
        logger.debug(
            f"stopReceiving() | calling pc.setLocalDescription() [answer:{answer}]"
        )
        await self.pc.setLocalDescription(answer)
        self._mapMidTransceiver.pop(localId, None)

    async def getReceiverStats(self, localId: str):
        self._assertRecvDirection()
        transceiver = self._mapMidTransceiver.get(localId)
        if not transceiver:
            raise Exception("associated RTCRtpTransceiver not found")
        return await transceiver.receiver.getStats()

    async def receiveDataChannel(
        self,
        sctpStreamParameters: SctpStreamParameters,
        label: Optional[str] = None,
        protocol: Optional[str] = None,
    ) -> HandlerReceiveDataChannelResult:
        options = HandlerReceiveDataChannelOptions(
            sctpStreamParameters=sctpStreamParameters, label=label, protocol=protocol
        )
        self._assertRecvDirection()
        logger.debug(f"[receiveDataChannel() [options:{options.sctpStreamParameters}]]")
        dataChannel = self.pc.createDataChannel(
            label=options.label,
            maxPacketLifeTime=options.sctpStreamParameters.maxPacketLifeTime,
            maxRetransmits=options.sctpStreamParameters.maxRetransmits,
            ordered=options.sctpStreamParameters.ordered,
            protocol=options.protocol,
            negotiated=True,
            id=options.sctpStreamParameters.streamId,
        )

        # If this is the first DataChannel we need to create the SDP offer with
        # m=application section.
        if not self._hasDataChannelMediaSection:
            self.remoteSdp.receiveSctpAssociation()
            offer: RTCSessionDescription = RTCSessionDescription(
                type="offer", sdp=self.remoteSdp.getSdp()
            )
            logger.debug(
                f"receiveDataChannel() | calling pc.setRemoteDescription() [offer:{offer}]"
            )
            await self.pc.setRemoteDescription(offer)
            answer = await self.pc.createAnswer()
            if not self._transportReady:
                localSdpDict = sdp_transform.parse(answer.sdp)
                await self._setupTransport(
                    localDtlsRole="client", localSdpDict=localSdpDict
                )
            logger.debug(
                f"receiveDataChannel() | calling pc.setRemoteDescription() [answer:{answer}]"
            )
            await self.pc.setLocalDescription(answer)
            self._hasDataChannelMediaSection = True
        return HandlerReceiveDataChannelResult(dataChannel=dataChannel)

    async def _setupTransport(self, localDtlsRole: DtlsRole, localSdpDict: dict = {}):
        if localSdpDict == {}:
            localSdpDict = sdp_transform.parse(self.pc.localDescription.sdp)
        # Get our local DTLS parameters.
        dtlsParameters: DtlsParameters = extractDtlsParameters(localSdpDict)
        # Set our DTLS role.
        dtlsParameters.role = localDtlsRole
        # Update the remote DTLS role in the SDP.
        self.remoteSdp.updateDtlsRole(
            "server" if localDtlsRole == "client" else "client"
        )
        # Need to tell the remote transport about our parameters.
        await self.emit_for_results("@connect", dtlsParameters)
        self._transportReady = True

    def _assertSendDirection(self):
        if self._direction != "send":
            raise Exception(
                'method can just be called for handlers with "send" direction'
            )

    def _assertRecvDirection(self):
        if self._direction != "recv":
            raise Exception(
                'method can just be called for handlers with "recv" direction'
            )
