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

import logging
from copy import deepcopy
from pyee.asyncio import AsyncIOEventEmitter
from aiortc import RTCIceServer, MediaStreamTrack
from .ortc import (
    canReceive,
    generateProbatorRtpParameters,
    ExtendedRtpCapabilities,
    validateAndNormalizeSctpStreamParameters,
)
from .errors import InvalidStateError, UnsupportedError
from .emitter import EnhancedEventEmitter
from .handlers.handler_interface import HandlerInterface
from .models.handler_interface import (
    HandlerSendResult,
    HandlerReceiveResult,
    SctpStreamParameters,
    HandlerSendDataChannelResult,
    HandlerReceiveDataChannelResult,
)
from .models.transport import (
    ConnectionState,
    IceParameters,
    InternalTransportOptions,
    DtlsParameters,
)
from .consumer import Consumer, ConsumerOptions
from .producer import Producer, ProducerOptions
from .data_consumer import DataConsumer, DataConsumerOptions
from .data_producer import DataProducer, DataProducerOptions
from .producer import ProducerCodecOptions
from .rtp_parameters import (
    RtpParameters,
    RtpCodecCapability,
    RtpEncodingParameters,
    MediaKind,
)


logger = logging.getLogger(__name__)


class Transport(EnhancedEventEmitter):
    def __init__(self, options: InternalTransportOptions, loop=None):
        super(Transport, self).__init__(loop=loop)

        logger.debug(f"constructor() [id:{options.id}, direction:{options.direction}]")

        # Closed flag.
        self._closed: bool = False
        # Whether we can produce audio/video based on computed extended RTP
        # capabilities.
        self._canProduceByKind: Dict[str, bool]
        # App custom data.
        self._appData: Optional[dict]
        # Transport connection state.
        self._connectionState: ConnectionState = "new"
        # Producers indexed by id
        self._producers: Dict[str, Producer] = {}
        # Consumers indexed by id.
        self._consumers: Dict[str, Consumer] = {}
        # DataProducers indexed by id
        self._dataProducers: Dict[str, DataProducer] = {}
        # DataConsumers indexed by id.
        self._dataConsumers: Dict[str, DataConsumer] = {}
        # Whether the Consumer for RTP probation has been created.
        self._probatorConsumerCreated: bool = False
        # Observer instance.
        self._observer: AsyncIOEventEmitter = AsyncIOEventEmitter()

        # Id.
        self._id: str = options.id
        # Direction.
        self._direction: Literal["send", "recv"] = options.direction
        # Extended RTP capabilities.
        extendedRtpCapabilities = options.extendedRtpCapabilities
        if extendedRtpCapabilities is None:
            raise TypeError("missing extendedRtpCapabilities")
        self._extendedRtpCapabilities: ExtendedRtpCapabilities = extendedRtpCapabilities
        self._canProduceByKind: Dict[str, bool] = options.canProduceByKind
        self._maxSctpMessageSize = (
            options.sctpParameters.maxMessageSize if options.sctpParameters else None
        )

        if options.additionalSettings:
            additionalSettings = deepcopy(options.additionalSettings)
            additionalSettings.pop("iceServers", None)
            additionalSettings.pop("iceTransportPolicy", None)
            additionalSettings.pop("bundlePolicy", None)
            additionalSettings.pop("rtcpMuxPolicy", None)
            additionalSettings.pop("sdpSemantics", None)
        else:
            additionalSettings = None

        self._handler: HandlerInterface = options.handlerFactory()

        self._handler.run(
            direction=options.direction,
            iceParameters=options.iceParameters,
            iceCandidates=options.iceCandidates,
            dtlsParameters=options.dtlsParameters,
            sctpParameters=options.sctpParameters,
            iceServers=options.iceServers,
            iceTransportPolicy=options.iceTransportPolicy,
            additionalSettings=additionalSettings,
            proprietaryConstraints=options.proprietaryConstraints,
            extendedRtpCapabilities=extendedRtpCapabilities,
        )

        self._appData = options.appData

        self._handleHandler()

    # Producer id.
    @property
    def id(self) -> str:
        return self._id

    # Whether the Producer is closed.
    @property
    def closed(self) -> bool:
        return self._closed

    # Transport direction.
    @property
    def direction(self) -> Literal["send", "recv"]:
        return self._direction

    # RTC handler instance.
    @property
    def handler(self) -> HandlerInterface:
        return self._handler

    # Connection state.
    @property
    def connectionState(self) -> ConnectionState:
        return self._connectionState

    # App custom data.
    @property
    def appData(self) -> Any:
        return self._appData

    # App custom data setter.
    @appData.setter
    def appData(self, value):
        self._appData = value

    # Observer.
    #
    # @emits close
    # @emits newproducer - (producer: Producer)
    # @emits newconsumer - (producer: Producer)
    # @emits newdataproducer - (dataProducer: DataProducer)
    # @emits newdataconsumer - (dataProducer: DataProducer)
    @property
    def observer(self) -> AsyncIOEventEmitter:
        return self._observer

    # Close the Transport.
    async def close(self):
        if self._closed:
            return

        logger.debug("Transport close()")

        self._closed = True
        self._connectionState = "closed"

        # Close the handler.
        await self._handler.close()

        # Close all Process.
        for process_dict in [
            self._producers,
            self._consumers,
            self._dataProducers,
            self._dataConsumers,
        ]:
            for process in process_dict.values():
                process.transportClosed()
            process_dict.clear()

        self._observer.emit("close")

    # Get associated Transport (RTCPeerConnection) stats.
    #
    # @returns {RTCStatsReport}
    async def getStats(self):
        if self._closed:
            raise InvalidStateError("closed")

        return await self._handler.getTransportStats()

    # Restart ICE connection.
    async def restartIce(self, iceParameters: IceParameters):
        if self._closed:
            raise InvalidStateError("closed")

        return await self._handler.restartIce(iceParameters)

    # Update ICE servers.
    async def updateIceServers(self, iceServers: List[RTCIceServer]):
        if self._closed:
            raise InvalidStateError("closed")

        return await self._handler.updateIceServers(iceServers)

    # Create a Producer.
    async def produce(
        self,
        track: Optional[MediaStreamTrack] = None,
        encodings: Optional[List[RtpEncodingParameters]] = None,
        codecOptions: Optional[ProducerCodecOptions] = None,
        codec: Optional[RtpCodecCapability] = None,
        streamId: Optional[str] = None,
        headerExtensionOptions: Optional[dict] = None,
        stopTracks: bool = True,
        disableTrackOnPause: bool = True,
        zeroRtpOnPause: bool = False,
        appData: Optional[Any] = None,
    ) -> Producer:
        options: ProducerOptions = ProducerOptions(
            track=track,
            encodings=encodings,
            codecOptions=codecOptions,
            codec=codec,
            streamId=streamId,
            headerExtensionOptions=headerExtensionOptions,
            stopTracks=stopTracks,
            disableTrackOnPause=disableTrackOnPause,
            zeroRtpOnPause=zeroRtpOnPause,
            appData=appData if appData is not None else {},
        )
        logger.debug(f"Transport produce() [track:{options.track}]")
        if self._closed:
            raise InvalidStateError("closed")
        elif not options.track:
            raise TypeError("missing track")
        elif self._direction != "send":
            raise UnsupportedError("not a sending Transport")
        elif not self._canProduceByKind.get(options.track.kind):
            raise UnsupportedError(f"cannot produce {options.track.kind}")
        elif options.track.readyState == "ended":
            raise InvalidStateError("track ended")
        elif len(self.listeners("connect")) == 0 and self._connectionState == "new":
            raise TypeError('no "connect" listener set into this transport')
        elif len(self.listeners("produce")) == 0:
            raise TypeError('no "produce" listener set into this transport')

        # NOTE: Mediasoup client enqueue command here.
        handlerSendResult: HandlerSendResult = await self._handler.send(
            track=options.track,
            encodings=options.encodings,
            codecOptions=options.codecOptions,
            codec=options.codec,
            streamId=options.streamId,
            headerExtensionOptions=options.headerExtensionOptions,
        )

        ids = await self.emit_for_results(
            "produce",
            options.track.kind,
            handlerSendResult.rtpParameters,
            options.appData,
        )

        producer = Producer(
            id=ids[0],
            localId=handlerSendResult.localId,
            rtpSender=handlerSendResult.rtpSender,
            track=options.track,
            rtpParameters=handlerSendResult.rtpParameters,
            stopTracks=options.stopTracks,
            disableTrackOnPause=options.disableTrackOnPause,
            zeroRtpOnPause=options.zeroRtpOnPause,
            appData=options.appData,
        )

        self._producers[producer.id] = producer
        self._handleProducer(producer)

        # Emit observer event.
        self._observer.emit("newproducer", producer)

        return producer

        # TODO: stop the given track if the command above failed due to closed Transport.

    async def consume(
        self,
        id: str,
        producerId: str,
        kind: MediaKind,
        rtpParameters: RtpParameters,
        streamId: Optional[str] = None,
        appData: Optional[dict] = None,
    ) -> Consumer:
        options: ConsumerOptions = ConsumerOptions(
            id=id,
            producerId=producerId,
            kind=kind,
            rtpParameters=rtpParameters,
            streamId=streamId,
            appData=appData if appData is not None else {},
        )
        logger.debug("Transport consume()")
        rtpParameters = options.rtpParameters.model_copy(deep=True)
        if self._closed:
            raise InvalidStateError("closed")
        elif self._direction != "recv":
            raise UnsupportedError("not a receiving Transport")
        elif len(self.listeners("connect")) == 0 and self._connectionState == "new":
            raise TypeError('no "connect" listener set into this transport')

        # NOTE: Mediasoup client enqueue command here.
        if not canReceive(
            rtpParameters=rtpParameters,
            extendedRtpCapabilities=self._extendedRtpCapabilities,
        ):
            raise UnsupportedError("cannot consume this Producer")

        handlerReceiveResult: HandlerReceiveResult = await self._handler.receive(
            trackId=options.id,
            kind=options.kind,
            rtpParameters=rtpParameters,
            streamId=options.streamId,
        )

        consumer: Consumer = Consumer(
            id=options.id,
            localId=handlerReceiveResult.localId,
            producerId=options.producerId,
            track=handlerReceiveResult.track,
            rtpParameters=rtpParameters,
            appData=options.appData,
        )

        self._consumers[consumer.id] = consumer
        self._handleConsumer(consumer)

        # If this is the first video Consumer and the Consumer for RTP probation
        # has not yet been created, create it now.
        if not self._probatorConsumerCreated and options.kind == "video":
            probatorRtpParameters = generateProbatorRtpParameters(
                consumer.rtpParameters
            )
            await self._handler.receive(
                trackId="probator", kind="video", rtpParameters=probatorRtpParameters
            )

            logger.debug("Transport consume() | Consumer for RTP probation created")

            self._probatorConsumerCreated = True

        self._observer.emit("newconsumer", consumer)

        return consumer

    # Create a DataProducer
    async def produceData(
        self,
        ordered: Optional[bool] = True,
        maxPacketLifeTime: Optional[int] = None,
        maxRetransmits: Optional[int] = None,
        label: Optional[str] = None,
        protocol: Optional[str] = None,
        appData: Optional[dict] = None,
    ) -> DataProducer:
        options: DataProducerOptions = DataProducerOptions(
            ordered=ordered,
            maxPacketLifeTime=maxPacketLifeTime,
            maxRetransmits=maxRetransmits,
            label=label,
            protocol=protocol,
            appData=appData if appData is not None else {},
        )
        logger.debug("Transport produceData()")
        if self._closed:
            raise InvalidStateError("closed")
        elif self._direction != "send":
            raise UnsupportedError("not a sending Transport")

        elif not self._maxSctpMessageSize:
            raise UnsupportedError("SCTP not enabled by remote Transport")

        elif len(self.listeners("connect")) == 0 and self._connectionState == "new":
            raise TypeError('no "connect" listener set into this transport')

        elif len(self.listeners("producedata")) == 0:
            raise TypeError('no "producedata" listener set into this transport')

        if options.maxPacketLifeTime or options.maxRetransmits:
            options.ordered = False

        # NOTE: Mediasoup client enqueue command here.
        handlerSendDataChannelResult: HandlerSendDataChannelResult = (
            await self._handler.sendDataChannel(
                ordered=options.ordered,
                maxPacketLifeTime=options.maxPacketLifeTime,
                maxRetransmits=options.maxRetransmits,
                label=options.label,
                protocol=options.protocol,
            )
        )

        ids = await self.emit_for_results(
            "producedata",
            sctpStreamParameters=handlerSendDataChannelResult.sctpStreamParameters,
            label=options.label,
            protocol=options.protocol,
            appData=options.appData,
        )

        dataProducer = DataProducer(
            id=ids[0],
            dataChannel=handlerSendDataChannelResult.dataChannel,
            sctpStreamParameters=handlerSendDataChannelResult.sctpStreamParameters,
            appData=options.appData,
        )

        self._dataProducers[dataProducer.id] = dataProducer
        self._handleDataProducer(dataProducer)

        # Emit observer event.
        self._observer.emit("newdataproducer", dataProducer)

        return dataProducer

    # Create a DataConsumer
    async def consumeData(
        self,
        id: str,
        dataProducerId: str,
        sctpStreamParameters: SctpStreamParameters,
        label: Optional[str] = None,
        protocol: Optional[str] = None,
        appData: Optional[dict] = None,
    ) -> DataConsumer:
        options: DataConsumerOptions = DataConsumerOptions(
            id=id,
            dataProducerId=dataProducerId,
            sctpStreamParameters=sctpStreamParameters,
            label=label,
            protocol=protocol,
            appData=appData if appData is not None else {},
        )
        logger.debug("Transport consumeData()")
        validateAndNormalizeSctpStreamParameters(options.sctpStreamParameters)
        if self._closed:
            raise InvalidStateError("closed")
        elif self._direction != "recv":
            raise UnsupportedError("not a receiving Transport")
        elif not self._maxSctpMessageSize:
            raise UnsupportedError("SCTP not enabled by remote Transport")
        elif len(self.listeners("connect")) == 0 and self._connectionState == "new":
            raise TypeError('no "connect" listener set into this transport')

        # NOTE: Mediasoup client enqueue command here.
        handlerReceiveDataChannelResult: HandlerReceiveDataChannelResult = (
            await self._handler.receiveDataChannel(
                sctpStreamParameters=options.sctpStreamParameters,
                label=options.label,
                protocol=options.protocol,
            )
        )

        dataConsumer: DataConsumer = DataConsumer(
            id=options.id,
            dataProducerId=options.dataProducerId,
            dataChannel=handlerReceiveDataChannelResult.dataChannel,
            sctpStreamParameters=options.sctpStreamParameters,
            appData=options.appData,
        )

        self._dataConsumers[dataConsumer.id] = dataConsumer
        self._handleDataConsumer(dataConsumer)

        # Emit observer event.
        self._observer.emit("newdataconsumer", dataConsumer)

        return dataConsumer

    def _handleHandler(self):
        handler = self._handler

        @handler.on("@connect")
        async def on_connect(dtlsParameters: DtlsParameters):
            if self._closed:
                raise InvalidStateError("closed")
            else:
                self.emit("connect", dtlsParameters)

        @handler.on("@connectionstatechange")
        def on_connectionstatechange(connectionState: ConnectionState):
            self._connectionState = connectionState
            if not self._closed:
                self.emit("connectionstatechange", connectionState)

    def _handleProducer(self, producer: Producer):
        @producer.on("@close")
        async def on_close():
            del self._producers[producer.id]
            if self._closed:
                return
            await self._handler.stopSending(producer.localId)

        @producer.on("@pause")
        async def on_pause():
            await self._handler.pauseSending(producer.localId)

        @producer.on("@resume")
        async def on_resume():
            await self._handler.resumeSending(producer.localId)

        @producer.on("@replacetrack")
        async def on_replacetrack(track):
            await self._handler.replaceTrack(producer.localId, track)

        @producer.on("@setmaxspatiallayer")
        async def on_setmaxspatiallayer(spatialLayer):
            await self._handler.setMaxSpatialLayer(producer.localId, spatialLayer)

        @producer.on("@setrtpencodingparameters")
        async def on_setrtpencodingparameters(params):
            await self._handler.setRtpEncodingParameters(producer.localId, params)

        @producer.on("@getstats")
        async def on_getstats():
            if self._closed:
                return InvalidStateError("closed")
            return await self._handler.getSenderStats(producer.localId)

    def _handleConsumer(self, consumer: Consumer):
        @consumer.on("@close")
        async def on_close():
            del self._consumers[consumer.id]
            if self._closed:
                return
            await self._handler.stopReceiving(consumer.localId)

        @consumer.on("@pause")
        async def on_pause():
            await self._handler.pauseReceiving(consumer.localId)

        @consumer.on("@resume")
        async def on_resume():
            await self._handler.resumeReceiving(consumer.localId)

        @consumer.on("@getstats")
        async def on_getstats():
            if self._closed:
                return InvalidStateError("closed")
            return await self._handler.getReceiverStats(consumer.localId)

    def _handleDataProducer(self, dataProducer: DataProducer):
        @dataProducer.on("@close")
        async def on_close():
            del self._dataProducers[dataProducer.id]

    def _handleDataConsumer(self, dataConsumer: DataConsumer):
        @dataConsumer.on("@close")
        async def on_close():
            del self._dataConsumers[dataConsumer.id]
