import asyncio
import json
import logging
import os
from fractions import Fraction

import aiohttp_cors
import aiortc.rtcrtpsender
import zmq
from aiohttp import web
from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription, RTCConfiguration, RTCRtpCodecCapability

from av import VideoFrame

import H264NV

ROOT = os.path.dirname(__file__)


class PlayerStreamTrackx(MediaStreamTrack):
    def __init__(self):
        super().__init__()
        context = zmq.Context()
        self.socket = context.socket(zmq.PULL)
        self.socket.bind("tcp://*:5555")
        self.kind = 'video'
        self.pts = 0

    async def recv(self):
        frame = VideoFrame.from_bytes(self.socket.recv(), 1920, 1080)
        frame.time_base = Fraction(1, 120)
        frame.pts = self.pts
        self.pts += 1
        return frame


pcs = set()
px = PlayerStreamTrackx()


async def offer(request):
    params = await request.json()
    offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
    pc = RTCPeerConnection(RTCConfiguration([]))
    # print(request.remote)
    pcs.add(pc)
    rc = pc.addTransceiver(px, 'sendonly')
    rc.setCodecPreferences([RTCRtpCodecCapability(mimeType='video/H264',
                                                  clockRate=90000,
                                                  channels=None,
                                                  parameters={
                                                      'level-asymmetry-allowed': '1',
                                                      'packetization-mode': '1',
                                                      'profile-level-id': '42e01f'
                                                  })])
    await pc.setRemoteDescription(offer)
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)
    return web.Response(
        content_type="application/json",
        text=json.dumps(
            {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}
        ),
    )


async def on_shutdown(app):
    # close peer connections
    coros = [pc.close() for pc in pcs]
    await asyncio.gather(*coros)
    pcs.clear()


if __name__ == "__main__":
    f0 = aiortc.RTCRtpSender.__init__


    def f1(*args, **kwargs):
        f0(*args, **kwargs)
        if not args[1] == 'audio':
            args[0]._RTCRtpSender__encoder = H264NV.H264NVENCEncoder()


    aiortc.RTCRtpSender.__init__ = f1
    logging.basicConfig(level=logging.INFO)
    app = web.Application()
    app.on_shutdown.append(on_shutdown)
    app.router.add_post("/offer", offer)

    cors = aiohttp_cors.setup(app, defaults={
        "*": aiohttp_cors.ResourceOptions(
            allow_credentials=True,
            expose_headers="*",
            allow_headers="*"
        )
    })

    for route in list(app.router.routes()):
        cors.add(route)

    web.run_app(
        app, access_log=None, host='0.0.0.0', port=8081
    )