#!/usr/bin/env python3
import asyncio
import json
import re

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import PlainTextResponse

app = FastAPI()
CODE_RE = re.compile(r"^\d{9}$")

# pending[code] = {"sender": ws|None, "receiver": ws|None, "event": asyncio.Event()}
PENDING = {}
PENDING_LOCK = asyncio.Lock()


@app.get("/")
async def root():
    return PlainTextResponse("ok")


async def register(code, role, ws):
    async with PENDING_LOCK:
        entry = PENDING.get(code)
        if entry is None:
            entry = {"sender": None, "receiver": None, "event": asyncio.Event()}
            PENDING[code] = entry
        if entry.get(role) is not None:
            return None, "role already connected"
        entry[role] = ws
        if entry.get("sender") is not None and entry.get("receiver") is not None:
            entry["event"].set()
        return entry, None


async def unregister(code, role, ws):
    async with PENDING_LOCK:
        entry = PENDING.get(code)
        if entry is None:
            return
        if entry.get(role) is ws:
            entry[role] = None
        if entry.get("sender") is None and entry.get("receiver") is None:
            PENDING.pop(code, None)


async def forward(src, dst):
    try:
        while True:
            msg = await src.receive()
            if msg.get("type") == "websocket.disconnect":
                break
            if msg.get("bytes") is not None:
                await dst.send_bytes(msg["bytes"])
            elif msg.get("text") is not None:
                await dst.send_text(msg["text"])
    except WebSocketDisconnect:
        pass
    except Exception:
        pass


@app.websocket("/ws")
async def ws_relay(ws: WebSocket):
    await ws.accept()
    code = None
    role = None
    entry = None
    try:
        raw = await ws.receive_text()
        try:
            payload = json.loads(raw)
        except json.JSONDecodeError:
            await ws.send_text(json.dumps({"error": "invalid json"}))
            return
        role = payload.get("role")
        code = payload.get("code")
        if role not in ("sender", "receiver"):
            await ws.send_text(json.dumps({"error": "invalid role"}))
            return
        if not isinstance(code, str) or not CODE_RE.match(code):
            await ws.send_text(json.dumps({"error": "invalid code"}))
            return
        entry, err = await register(code, role, ws)
        if err:
            await ws.send_text(json.dumps({"error": err}))
            return
        await ws.send_text(json.dumps({"status": "waiting"}))
        await entry["event"].wait()
        other = entry["receiver"] if role == "sender" else entry["sender"]
        if other is None:
            await ws.send_text(json.dumps({"error": "peer missing"}))
            return
        await ws.send_text(json.dumps({"status": "paired"}))
        await other.send_text(json.dumps({"status": "paired"}))
        task_a = asyncio.create_task(forward(ws, other))
        task_b = asyncio.create_task(forward(other, ws))
        done, pending = await asyncio.wait(
            {task_a, task_b}, return_when=asyncio.FIRST_COMPLETED
        )
        for task in pending:
            task.cancel()
    finally:
        if code and role:
            await unregister(code, role, ws)
        try:
            await ws.close()
        except Exception:
            pass