xenia-pos-local/local_backend/services/ws_bus.py

"""
WebSocket Event Bus — replaces sse_bus.py for real-time communication.

Drop-in replacement: all routers continue calling broadcast_sync() unchanged.
The WS endpoint uses connect()/disconnect()/replay_missed() to manage clients.

Protocol (JSON frames):
  Server → Client:  { "seq": 123, "type": "order_updated", "data": { ... } }
  Client → Server:  { "cursor": 120 }   (sent immediately after connect)
  Server → Client:  { "type": "ping" }  (every 25s keepalive)
  Client → Server:  { "type": "pong" }  (optional, ignored if not sent)

On connect the client sends its last known seq. The server replays everything
it has stored since that seq, then switches to live streaming. Events are stored
in the sync_events SQLite table (written here, pruned on startup).
"""

import asyncio
import json
import logging
from typing import Dict, Set
from datetime import datetime, timedelta, timezone

from fastapi import WebSocket

logger = logging.getLogger(__name__)

# ── Event loop (captured once at startup) ────────────────────────────────────

_main_loop: asyncio.AbstractEventLoop | None = None


def init_loop(loop: asyncio.AbstractEventLoop) -> None:
    global _main_loop
    _main_loop = loop


# ── Connected clients: user_id → set of asyncio.Queue ────────────────────────

_connections: Dict[int, Set[asyncio.Queue]] = {}


async def connect(user_id: int) -> asyncio.Queue:
    q: asyncio.Queue = asyncio.Queue(maxsize=512)
    if user_id not in _connections:
        _connections[user_id] = set()
    _connections[user_id].add(q)
    return q


async def disconnect(user_id: int, q: asyncio.Queue) -> None:
    if user_id in _connections:
        _connections[user_id].discard(q)
        if not _connections[user_id]:
            del _connections[user_id]


# ── Persistence helpers ───────────────────────────────────────────────────────

def _get_db():
    from database import SessionLocal
    return SessionLocal()


def _persist_event(event_type: str, data: dict, user_ids: list[int] | None) -> int:
    """Write event to sync_events table, return the new seq_id."""
    from sqlalchemy import text
    db = _get_db()
    try:
        result = db.execute(
            text(
                "INSERT INTO sync_events (event_type, payload, target_user_ids, created_at) "
                "VALUES (:et, :payload, :uids, :now)"
            ),
            {
                "et": event_type,
                "payload": json.dumps(data),
                "uids": json.dumps(user_ids) if user_ids is not None else None,
                "now": datetime.now(timezone.utc).isoformat(),
            },
        )
        db.commit()
        return result.lastrowid
    finally:
        db.close()


def get_events_since(cursor: int, user_id: int) -> list[dict]:
    """Return all events with seq_id > cursor that are visible to user_id."""
    from sqlalchemy import text
    db = _get_db()
    try:
        rows = db.execute(
            text(
                "SELECT id, event_type, payload, target_user_ids FROM sync_events "
                "WHERE id > :cursor ORDER BY id ASC LIMIT 500"
            ),
            {"cursor": cursor},
        ).fetchall()
    finally:
        db.close()

    events = []
    for row in rows:
        target_user_ids = json.loads(row[3]) if row[3] else None
        if target_user_ids is None or user_id in target_user_ids:
            events.append({
                "seq": row[0],
                "type": row[1],
                "data": json.loads(row[2]),
            })
    return events


def prune_old_events(hours: int = 24) -> int:
    """Delete events older than `hours`. Called on startup."""
    from sqlalchemy import text
    db = _get_db()
    try:
        cutoff = (datetime.now(timezone.utc) - timedelta(hours=hours)).isoformat()
        result = db.execute(
            text("DELETE FROM sync_events WHERE created_at < :cutoff"),
            {"cutoff": cutoff},
        )
        db.commit()
        return result.rowcount
    finally:
        db.close()


# ── Broadcast ─────────────────────────────────────────────────────────────────

def broadcast_sync(event_type: str, data: dict, *, user_ids: list[int] | None = None) -> None:
    """
    Called from sync route thread-pool workers.
    Persists the event to DB (so reconnecting clients can replay it),
    then schedules an async push to all currently connected sockets.
    """
    try:
        seq_id = _persist_event(event_type, data, user_ids)
    except Exception:
        logger.exception("ws_bus: failed to persist event %s", event_type)
        seq_id = 0

    if _main_loop is None:
        return
    _main_loop.call_soon_threadsafe(
        _main_loop.create_task,
        _broadcast_live(seq_id, event_type, data, user_ids),
    )


async def _broadcast_live(
    seq_id: int,
    event_type: str,
    data: dict,
    user_ids: list[int] | None,
) -> None:
    frame = json.dumps({"seq": seq_id, "type": event_type, "data": data})
    targets = (
        {uid: qs for uid, qs in _connections.items() if uid in user_ids}
        if user_ids is not None
        else dict(_connections)
    )
    for qs in targets.values():
        for q in list(qs):
            try:
                q.put_nowait(frame)
            except asyncio.QueueFull:
                pass  # slow client — drop live frame; they'll replay on reconnect