WebSocket Architecture¶

from enum import Enum


class RoomType(str, Enum):
    WORKBOOK = "workbook"
    USER = "user"
    TENANT = "tenant"


class Room:
    """Represents a WebSocket room with type-safe naming."""

    def __init__(self, room_type: RoomType, resource_id: str):
        self.room_type = room_type
        self.resource_id = resource_id

    @property
    def channel_name(self) -> str:
        """Valkey pub/sub channel name for this room."""
        return f"ws:{self.room_type.value}:{self.resource_id}"

    @classmethod
    def workbook(cls, workbook_id: str) -> "Room":
        return cls(RoomType.WORKBOOK, workbook_id)

    @classmethod
    def user(cls, user_id: str) -> "Room":
        return cls(RoomType.USER, user_id)

    @classmethod
    def tenant(cls, tenant_id: str) -> "Room":
        return cls(RoomType.TENANT, tenant_id)

    def __str__(self) -> str:
        return self.channel_name

from dataclasses import dataclass, field
from enum import Enum
import time
import uuid


class EventType(str, Enum):
    # Calculation events
    CALCULATION_PROGRESS = "calculation_progress"
    CALCULATION_COMPLETE = "calculation_complete"
    CALCULATION_FAILED = "calculation_failed"

    # Workbook events
    WORKBOOK_UPDATED = "workbook_updated"
    WORKBOOK_DELETED = "workbook_deleted"

    # Export events
    EXPORT_PROGRESS = "export_progress"
    EXPORT_COMPLETE = "export_complete"
    EXPORT_FAILED = "export_failed"

    # Scenario events
    SCENARIO_COMPLETE = "scenario_complete"
    SCENARIO_FAILED = "scenario_failed"

    # Reconciliation events
    RECON_COMPLETE = "reconciliation_complete"
    RECON_DISCREPANCIES = "reconciliation_discrepancies"

    # Notification events
    NOTIFICATION = "notification"

    # System events
    SYSTEM_ANNOUNCEMENT = "system_announcement"
    MAINTENANCE_NOTICE = "maintenance_notice"

    # Connection events
    SUBSCRIBED = "subscribed"
    UNSUBSCRIBED = "unsubscribed"
    ERROR = "error"


@dataclass
class WebSocketEvent:
    """Standard event envelope for all WebSocket messages."""

    type: EventType
    payload: dict
    version: str = "1.0"
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: float = field(default_factory=time.time)
    room: str = ""

    def to_json(self) -> str:
        """Serialize event to JSON string for transmission."""
        import json
        return json.dumps({
            "type": self.type.value,
            "version": self.version,
            "event_id": self.event_id,
            "timestamp": self.timestamp,
            "room": self.room,
            "payload": self.payload,
        })

import asyncio
from collections import defaultdict
from dataclasses import dataclass, field
from fastapi import WebSocket


@dataclass
class ConnectionState:
    """Tracks state for a single WebSocket connection."""

    websocket: WebSocket
    user_id: str
    tenant_id: str
    roles: list[str]
    rooms: set[str] = field(default_factory=set)
    connected_at: float = 0.0


class ConnectionManager:
    """Manages WebSocket connections, room subscriptions,
    and event delivery.

    Thread-safe for use with asyncio. Uses in-memory structures
    for connection tracking (not persisted -- connections are
    ephemeral by nature).
    """

    def __init__(self):
        self._connections: dict[str, ConnectionState] = {}
        self._rooms: dict[str, set[str]] = defaultdict(set)
        self._lock = asyncio.Lock()

    async def register(
        self,
        connection_id: str,
        websocket: WebSocket,
        user_id: str,
        tenant_id: str,
        roles: list[str],
    ) -> ConnectionState:
        """Register a new authenticated connection.

        Automatically subscribes to user and tenant rooms.
        """
        async with self._lock:
            state = ConnectionState(
                websocket=websocket,
                user_id=user_id,
                tenant_id=tenant_id,
                roles=roles,
                connected_at=time.time(),
            )
            self._connections[connection_id] = state

            # Auto-subscribe to user and tenant rooms
            user_room = f"ws:user:{user_id}"
            tenant_room = f"ws:tenant:{tenant_id}"

            state.rooms.add(user_room)
            state.rooms.add(tenant_room)
            self._rooms[user_room].add(connection_id)
            self._rooms[tenant_room].add(connection_id)

            return state

    async def unregister(self, connection_id: str) -> None:
        """Remove a connection and clean up all room subscriptions."""
        async with self._lock:
            state = self._connections.pop(connection_id, None)
            if state:
                for room in state.rooms:
                    self._rooms[room].discard(connection_id)
                    if not self._rooms[room]:
                        del self._rooms[room]

    async def join_room(
        self, connection_id: str, room: str
    ) -> bool:
        """Subscribe a connection to a room.

        Returns False if the connection does not have permission
        to join the requested room (e.g., workbook belongs to
        a different tenant).
        """
        async with self._lock:
            state = self._connections.get(connection_id)
            if not state:
                return False

            # Permission check: workbook rooms require tenant match
            if room.startswith("ws:workbook:"):
                if not await self._verify_workbook_access(
                    state.tenant_id, room
                ):
                    return False

            state.rooms.add(room)
            self._rooms[room].add(connection_id)
            return True

    async def leave_room(
        self, connection_id: str, room: str
    ) -> None:
        """Unsubscribe a connection from a room."""
        async with self._lock:
            state = self._connections.get(connection_id)
            if state:
                state.rooms.discard(room)
                self._rooms[room].discard(connection_id)

    async def broadcast_to_room(
        self, room: str, event: WebSocketEvent
    ) -> int:
        """Send an event to all connections in a room.

        Returns the number of connections that received the event.
        Failed sends (disconnected clients) are cleaned up.
        """
        event.room = room
        message = event.to_json()
        delivered = 0
        failed = []

        connection_ids = self._rooms.get(room, set()).copy()

        for conn_id in connection_ids:
            state = self._connections.get(conn_id)
            if state:
                try:
                    await state.websocket.send_text(message)
                    delivered += 1
                except Exception:
                    failed.append(conn_id)

        # Clean up failed connections
        for conn_id in failed:
            await self.unregister(conn_id)

        return delivered

    @property
    def active_connections(self) -> int:
        """Total number of active connections."""
        return len(self._connections)

    @property
    def active_rooms(self) -> int:
        """Total number of rooms with at least one subscriber."""
        return len(self._rooms)

import json
from valkey import Valkey


class EventPublisher:
    """Publishes WebSocket events from Celery workers via Valkey pub/sub.

    Workers do not maintain WebSocket connections directly.
    Instead, they publish events to Valkey channels, and the
    API process's subscriber forwards them to connected clients.
    """

    def __init__(self, valkey_client: Valkey):
        self._valkey = valkey_client

    def publish(self, room: str, event: WebSocketEvent) -> int:
        """Publish an event to a Valkey channel.

        Returns the number of subscribers that received the message
        (Valkey pub/sub subscriber count, not WebSocket clients).
        """
        event.room = room
        return self._valkey.publish(room, event.to_json())

    def publish_calculation_progress(
        self,
        workbook_id: str,
        job_id: str,
        progress_pct: int,
        current_sheet: str,
        cells_processed: int,
        cells_total: int,
    ) -> None:
        """Convenience method for calculation progress events."""
        event = WebSocketEvent(
            type=EventType.CALCULATION_PROGRESS,
            payload={
                "workbook_id": workbook_id,
                "job_id": job_id,
                "progress_pct": progress_pct,
                "current_sheet": current_sheet,
                "cells_processed": cells_processed,
                "cells_total": cells_total,
                "eta_seconds": self._estimate_eta(
                    progress_pct, cells_processed, cells_total
                ),
            },
        )
        room = Room.workbook(workbook_id).channel_name
        self.publish(room, event)

import asyncio
import json


class EventSubscriber:
    """Subscribes to Valkey pub/sub channels and forwards events
    to the ConnectionManager for WebSocket delivery.

    Runs as a background asyncio task within the FastAPI process.
    """

    def __init__(
        self,
        valkey_client,
        connection_manager: ConnectionManager,
    ):
        self._valkey = valkey_client
        self._manager = connection_manager
        self._pubsub = None
        self._running = False

    async def start(self) -> None:
        """Start the pub/sub listener as a background task."""
        self._pubsub = self._valkey.pubsub()
        self._running = True
        asyncio.create_task(self._listen())

    async def stop(self) -> None:
        """Stop the pub/sub listener and clean up."""
        self._running = False
        if self._pubsub:
            await self._pubsub.unsubscribe()
            await self._pubsub.close()

    async def subscribe_to_room(self, room: str) -> None:
        """Subscribe to a Valkey channel for a room."""
        if self._pubsub:
            await self._pubsub.subscribe(room)

    async def unsubscribe_from_room(self, room: str) -> None:
        """Unsubscribe from a Valkey channel."""
        if self._pubsub:
            await self._pubsub.unsubscribe(room)

    async def _listen(self) -> None:
        """Main listener loop -- forwards Valkey messages
        to WebSocket clients via ConnectionManager.
        """
        while self._running:
            try:
                message = await self._pubsub.get_message(
                    ignore_subscribe_messages=True,
                    timeout=1.0,
                )
                if message and message["type"] == "message":
                    channel = message["channel"]
                    event_data = json.loads(message["data"])
                    event = WebSocketEvent(**event_data)

                    await self._manager.broadcast_to_room(
                        channel, event
                    )
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Pub/sub listener error: {e}")
                await asyncio.sleep(1)  # Back off on error

from fastapi import WebSocket, status
from src.core.security import verify_jwt_token


async def authenticate_websocket(
    websocket: WebSocket,
) -> dict | None:
    """Authenticate a WebSocket connection via JWT.

    The token can be provided in two ways:
    1. Authorization header: "Bearer <token>"
    2. Query parameter: "?token=<token>"

    Returns user context dict or None if authentication fails.
    """
    # Try Authorization header first
    token = None
    auth_header = websocket.headers.get("authorization")
    if auth_header and auth_header.startswith("Bearer "):
        token = auth_header[7:]

    # Fall back to query parameter
    if not token:
        token = websocket.query_params.get("token")

    if not token:
        await websocket.close(
            code=4001, reason="Missing authentication token"
        )
        return None

    try:
        payload = verify_jwt_token(token)
        return {
            "user_id": payload["sub"],
            "tenant_id": payload["tenant_id"],
            "roles": payload.get("roles", []),
        }
    except Exception:
        await websocket.close(
            code=4001, reason="Invalid or expired token"
        )
        return None

async def _send_with_timeout(
    self, websocket: WebSocket, message: str, timeout: float = 5.0
) -> bool:
    """Send a message with timeout to detect slow clients."""
    try:
        await asyncio.wait_for(
            websocket.send_text(message), timeout=timeout
        )
        return True
    except asyncio.TimeoutError:
        logger.warning("Client send timeout, disconnecting")
        return False
    except Exception:
        return False