import argparse
import asyncio
import logging
import random
import string
import time

import aiohttp

from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack
from aiortc.contrib.media import MediaPlayer


def transaction_id():
    return ''.join(random.choice(string.ascii_letters) for x in range(12))


class JanusPlugin:
    def __init__(self, session, url):
        self._queue = asyncio.Queue()
        self._session = session
        self._url = url

    async def send(self, payload):
        message = {
            'janus': 'message',
            'transaction': transaction_id()
        }
        message.update(payload)
        async with self._session._http.post(self._url, json=message) as response:
            data = await response.json()
            assert data['janus'] == 'ack'

        response = await self._queue.get()
        assert response['transaction'] == message['transaction']
        return response


class JanusSession:
    def __init__(self, url):
        self._http = None
        self._poll_task = None
        self._plugins = {}
        self._root_url = url
        self._session_url = None

    async def attach(self, plugin):
        message = {
            'janus': 'attach',
            'plugin': plugin,
            'transaction': transaction_id(),
        }
        async with self._http.post(self._session_url, json=message) as response:
            data = await response.json()
            assert data['janus'] == 'success'
            plugin_id = data['data']['id']
            plugin = JanusPlugin(self, self._session_url + '/' + str(plugin_id))
            self._plugins[plugin_id] = plugin
            return plugin

    async def create(self):
        self._http = aiohttp.ClientSession()
        message = {
            'janus': 'create',
            'transaction': transaction_id(),
        }
        async with self._http.post(self._root_url, json=message) as response:
            data = await response.json()
            assert data['janus'] == 'success'
            session_id = data['data']['id']
            self._session_url = self._root_url + '/' + str(session_id)

        self._poll_task = asyncio.ensure_future(self._poll())

    async def destroy(self):
        if self._poll_task:
            self._poll_task.cancel()
            self._poll_task = None

        if self._session_url:
            message = {
                'janus': 'destroy',
                'transaction': transaction_id(),
            }
            async with self._http.post(self._session_url, json=message) as response:
                data = await response.json()
                assert data['janus'] == 'success'
            self._session_url = None

        if self._http:
            await self._http.close()
            self._http = None

    async def _poll(self):
        while True:
            params = {
                'maxev': 1,
                'rid': int(time.time() * 1000)
            }
            async with self._http.get(self._session_url, params=params) as response:
                data = await response.json()
                if data['janus'] == 'event':
                    plugin = self._plugins.get(data['sender'], None)
                    if plugin:
                        await plugin._queue.put(data)
                    else:
                        print(data)


async def run(pc, player, room, session):
    await session.create()

    # configure media
    media = {'audio': False, 'video': True}
    if player and player.audio:
        pc.addTrack(player.audio)
        media['audio'] = True

    if player and player.video:
        pc.addTrack(player.video)
    else:
        pc.addTrack(VideoStreamTrack())

    # join video room
    plugin = await session.attach('janus.plugin.videoroom')
    await plugin.send({
        'body': {
            'display': 'aiortc',
            'ptype': 'publisher',
            'request': 'join',
            'room': room,
        }
    })

    # send offer
    await pc.setLocalDescription(await pc.createOffer())
    request = {'request': 'configure'}
    request.update(media)
    response = await plugin.send({
        'body': request,
        'jsep': {
            'sdp': pc.localDescription.sdp,
            'trickle': False,
            'type': pc.localDescription.type
        }
    })

    # apply answer
    answer = RTCSessionDescription(
        sdp=response['jsep']['sdp'],
        type=response['jsep']['type'])
    await pc.setRemoteDescription(answer)

    # exchange media for 10 minutes
    print('Exchanging media')
    await asyncio.sleep(600)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Janus')
    parser.add_argument('url', help='Janus root URL, e.g. http://localhost:8088/janus')
    parser.add_argument('--room', type=int, default=1234,
                        help='The video room ID to join (default: 1234).'),
    parser.add_argument('--play-from', help='Read the media from a file and sent it.'),
    parser.add_argument('--verbose', '-v', action='count')
    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)

    # create signaling and peer connection
    session = JanusSession(args.url)
    pc = RTCPeerConnection()

    # create media source
    if args.play_from:
        player = MediaPlayer(args.play_from)
    else:
        player = None

    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(run(
            pc=pc,
            player=player,
            room=args.room,
            session=session))
    except KeyboardInterrupt:
        pass
    finally:
        loop.run_until_complete(pc.close())
        loop.run_until_complete(session.destroy())
