LMS/backend/app/routers/ws.py

"""
WebSocket collaboration endpoint.

URL: ws://.../ws/pdf/{pdf_id}?token=<JWT>

Each PDF has its own "room". When a client connects it broadcasts a `presence`
event to the room. All annotation mutations are forwarded verbatim to every
other client in the room.

Message schema (JSON):
  → client sends:
      { "type": "object_add",    "payload": <fabric object JSON> }
      { "type": "object_remove", "payload": { "obj_id": "<id>" } }
      { "type": "clear" }
      { "type": "cursor",        "payload": { "x": 0, "y": 0, "page": 1 } }
      { "type": "ping" }

  ← server sends to ALL others in room:
      same messages, with "user_id" / "username" / "color" injected

  ← server sends to ALL (including sender) on join/leave:
      { "type": "presence", "users": [ { "user_id", "username", "color" } ] }
"""

import asyncio
import json
import logging
from typing import Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status
from jose import JWTError

from ..security import decode_access_token
from ..database import get_db
from ..models import User

logger = logging.getLogger(__name__)
router = APIRouter()

# ── Deterministic colour per user (hue based on user_id) ─────────────────────

def _user_color(user_id: int) -> str:
    hue = (user_id * 61) % 360          # spread nicely around the wheel
    return f"hsl({hue},70%,50%)"


# ── Room manager ──────────────────────────────────────────────────────────────

class _Room:
    def __init__(self) -> None:
        # websocket → user info dict
        self._clients: dict[WebSocket, dict[str, Any]] = {}
        self._lock = asyncio.Lock()

    async def join(self, ws: WebSocket, user_info: dict[str, Any]) -> None:
        async with self._lock:
            self._clients[ws] = user_info
        await self._broadcast_presence()

    async def leave(self, ws: WebSocket) -> None:
        async with self._lock:
            self._clients.pop(ws, None)
        await self._broadcast_presence()

    @property
    def user_list(self) -> list[dict[str, Any]]:
        """Deduplicated by user_id — same user may have multiple WS connections."""
        seen: set[int] = set()
        result: list[dict[str, Any]] = []
        for info in self._clients.values():
            uid = info["user_id"]
            if uid not in seen:
                seen.add(uid)
                result.append(info)
        return result

    async def broadcast(self, message: dict[str, Any], exclude: WebSocket | None = None) -> None:
        """Send message to every client except `exclude`."""
        async with self._lock:
            targets = [ws for ws in self._clients if ws is not exclude]
        for ws in targets:
            try:
                await ws.send_json(message)
            except Exception:
                pass

    async def broadcast_all(self, message: dict[str, Any]) -> None:
        """Send message to every client including sender."""
        await self.broadcast(message, exclude=None)

    async def _broadcast_presence(self) -> None:
        await self.broadcast_all({"type": "presence", "users": self.user_list})


class _RoomManager:
    def __init__(self) -> None:
        self._rooms: dict[int, _Room] = {}
        self._lock = asyncio.Lock()

    async def get_or_create(self, pdf_id: int) -> _Room:
        async with self._lock:
            if pdf_id not in self._rooms:
                self._rooms[pdf_id] = _Room()
            return self._rooms[pdf_id]

    async def cleanup(self, pdf_id: int) -> None:
        async with self._lock:
            room = self._rooms.get(pdf_id)
            if room and not room._clients:
                del self._rooms[pdf_id]


manager = _RoomManager()

# ── WebSocket endpoint ────────────────────────────────────────────────────────

@router.websocket("/ws/pdf/{pdf_id}")
async def pdf_collaboration(websocket: WebSocket, pdf_id: int):
    # ── Auth: JWT passed as query param (cookies aren't reliably sent on WS) ──
    token = websocket.query_params.get("token")
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    try:
        user_id = decode_access_token(token)
    except JWTError:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    # ── Fetch username from DB ────────────────────────────────────────────────
    db = next(get_db())
    try:
        user: User | None = db.get(User, user_id)
    finally:
        db.close()

    if user is None:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    color = _user_color(user_id)
    user_info = {"user_id": user_id, "username": user.username, "color": color}

    await websocket.accept()
    room = await manager.get_or_create(pdf_id)
    await room.join(websocket, user_info)

    try:
        while True:
            raw = await websocket.receive_text()
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                continue

            msg_type = msg.get("type")

            if msg_type == "ping":
                await websocket.send_json({"type": "pong"})
                continue

            # Inject sender identity and forward to all other clients
            if msg_type in {"object_add", "object_remove", "clear", "cursor", "color_sync"}:
                msg.update(user_info)
                await room.broadcast(msg, exclude=websocket)

    except WebSocketDisconnect:
        pass
    finally:
        await room.leave(websocket)
        await manager.cleanup(pdf_id)