#!/usr/bin/env python3
"""Kokkoja relay.

Passes messages between the participants of a session.  It stores nothing:
no accounts, no pages, no notebooks.  Message bodies are encrypted by the
clients with a key the relay never sees, so the relay cannot read them.

The protocol is described in notes/protocol.md.  In short:

* A client connects to ``/room/<room id>``.  The relay gives the connection
  a peer number, unique within the room, and sends ``welcome``.

* A binary message from a client begins with a 4-byte destination peer
  number (0 means everyone else in the room).  The relay replaces those four
  bytes with the sender's peer number and forwards the message.

* The relay sends ``join`` and ``leave`` as peers come and go.

Requires the ``websockets`` library, version 13 or later.
"""

import argparse
import asyncio
import http
import json
import logging
import re
import signal

from websockets.asyncio.server import broadcast, serve
from websockets.exceptions import ConnectionClosed

ROOM_PATH = re.compile(r"^/room/([A-Za-z0-9_-]{22})$")

MAX_MESSAGE_BYTES = 1 << 20     # Clients split anything larger.
MAX_PEERS_PER_ROOM = 32
MAX_ROOMS = 500

# Close codes in the range reserved for applications.
CLOSE_ROOM_FULL = 4001
CLOSE_RELAY_FULL = 4002
CLOSE_BAD_MESSAGE = 4003

log = logging.getLogger("kokkoja_relay")


class Room:
    def __init__(self):
        self.peers = {}         # peer number -> connection
        self.next_peer = 1      # Peer numbers are never reused in a room.


rooms = {}                      # room id -> Room


def check_path(connection, request):
    """Refuse the WebSocket handshake unless the path names a room."""
    if not ROOM_PATH.match(request.path):
        return connection.respond(http.HTTPStatus.NOT_FOUND, "Not a Kokkoja room\n")
    return None


def control(message_type, **fields):
    return json.dumps({"type": message_type, **fields})


async def handle(connection):
    room_id = ROOM_PATH.match(connection.request.path).group(1)
    room = rooms.get(room_id)
    if room is None:
        if len(rooms) >= MAX_ROOMS:
            await connection.close(CLOSE_RELAY_FULL, "relay full")
            return
        room = rooms[room_id] = Room()
    if len(room.peers) >= MAX_PEERS_PER_ROOM:
        await connection.close(CLOSE_ROOM_FULL, "room full")
        return

    peer = room.next_peer
    room.next_peer += 1
    others = sorted(room.peers)
    room.peers[peer] = connection
    sender = peer.to_bytes(4, "big")
    log.info("room %s..: peer %d joined (%d present)", room_id[:4], peer, len(room.peers))

    try:
        await connection.send(control("welcome", peer=peer, peers=others))
        broadcast([room.peers[p] for p in others], control("join", peer=peer))
        async for message in connection:
            if not isinstance(message, bytes) or len(message) < 4:
                await connection.close(CLOSE_BAD_MESSAGE, "bad message")
                break
            destination = int.from_bytes(message[:4], "big")
            if destination == 0:
                targets = [c for p, c in room.peers.items() if p != peer]
            else:
                target = room.peers.get(destination)
                targets = [target] if target is not None and destination != peer else []
            broadcast(targets, sender + message[4:])
    except ConnectionClosed:
        pass
    finally:
        del room.peers[peer]
        if room.peers:
            broadcast(list(room.peers.values()), control("leave", peer=peer))
        else:
            del rooms[room_id]
        log.info("room %s..: peer %d left (%d present)", room_id[:4], peer, len(room.peers))


async def run(host, port):
    stop = asyncio.get_running_loop().create_future()

    def stop_once():
        if not stop.done():
            stop.set_result(None)

    for signal_number in (signal.SIGINT, signal.SIGTERM):
        asyncio.get_running_loop().add_signal_handler(signal_number, stop_once)
    async with serve(handle, host, port, process_request=check_path,
                     max_size=MAX_MESSAGE_BYTES, ping_interval=20, ping_timeout=20):
        log.info("relay listening on %s port %d", host, port)
        await stop


def main():
    parser = argparse.ArgumentParser(description="Kokkoja relay")
    parser.add_argument("--host", default="127.0.0.1",
                        help="address to listen on (default 127.0.0.1; 0.0.0.0 for all)")
    parser.add_argument("--port", type=int, default=8770)
    arguments = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
    # The library reports every connection as well; our own lines suffice.
    logging.getLogger("websockets").setLevel(logging.WARNING)
    asyncio.run(run(arguments.host, arguments.port))


if __name__ == "__main__":
    main()