led2/lightsync/api/ws.py

"""WebSocket hub — browser-driven tick/beat protocol (Phase 02-03 redesign).

Browser sends position ticks at ~10Hz; server looks up the next beat and
sends a beat notification back when one is near.

Protocol (browser → server):
    {"type": "tick",  "position": 42.3}          — audio playback position
    {"type": "play"}                               — playback started (log only for now)
    {"type": "pause"}                              — playback paused (log only for now)
    {"type": "seek",  "position": 42.3}           — seek event (log only for now)

Protocol (server → browser):
    {"type": "beat",  "position": 42.3, "tempo": 120.0}  — beat detected near position
    {"type": "error", "message": "..."}                   — error response
"""
from __future__ import annotations

import bisect
import json
import logging
from fastapi import APIRouter, WebSocket, WebSocketDisconnect

logger = logging.getLogger(__name__)

router = APIRouter()

# Beat detection tolerance window: consider a beat "hit" if audio position
# is within this many seconds of a known beat timestamp.
BEAT_WINDOW_SEC = 0.05  # 50ms


class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        self.active_connections.append(ws)

    def disconnect(self, ws: WebSocket) -> None:
        if ws in self.active_connections:
            self.active_connections.remove(ws)

    async def broadcast(self, message: dict) -> None:
        dead = []
        for connection in self.active_connections:
            try:
                await connection.send_text(json.dumps(message))
            except Exception:
                dead.append(connection)
        for d in dead:
            if d in self.active_connections:
                self.active_connections.remove(d)


manager = ConnectionManager()


def _find_nearest_beat(beats: list[float], position: float) -> float | None:
    """Return the nearest beat timestamp to position, or None if beats is empty.

    Uses bisect for O(log n) lookup.
    """
    if not beats:
        return None
    idx = bisect.bisect_left(beats, position)
    candidates = []
    if idx < len(beats):
        candidates.append(beats[idx])
    if idx > 0:
        candidates.append(beats[idx - 1])
    return min(candidates, key=lambda b: abs(b - position))


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    # Per-connection state: which file's beats are we tracking
    current_file: str | None = None
    beats: list[float] = []
    tempo: float = 0.0
    last_beat_reported: float | None = None  # avoid re-reporting same beat

    try:
        while True:
            data = await websocket.receive_text()
            try:
                msg = json.loads(data)
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({"type": "error", "message": "invalid JSON"}))
                continue

            msg_type = msg.get("type")

            if msg_type == "tick":
                position = float(msg.get("position", 0.0))

                # Lazily load beat data from app state when available
                if not beats:
                    beats_cache = getattr(websocket.app.state, "beats", None) or {}
                    # Use the most recently cached file if we don't have one yet
                    if beats_cache:
                        # Pick the last uploaded file (most likely what's playing)
                        latest_path = list(beats_cache.keys())[-1]
                        beat_data = beats_cache[latest_path]
                        current_file = latest_path
                        beats = sorted(beat_data.get("beats", []))
                        tempo = float(beat_data.get("tempo", 0.0))

                if beats:
                    nearest = _find_nearest_beat(beats, position)
                    if nearest is not None and abs(nearest - position) <= BEAT_WINDOW_SEC:
                        # Only send each beat once per playback pass
                        if last_beat_reported != nearest:
                            last_beat_reported = nearest
                            await websocket.send_text(json.dumps({
                                "type": "beat",
                                "position": nearest,
                                "tempo": tempo,
                            }))

            elif msg_type == "play":
                logger.info("[ws] play event at position=%.3f", float(msg.get("position", 0)))
                last_beat_reported = None  # reset beat tracking on play

            elif msg_type == "pause":
                logger.info("[ws] pause event at position=%.3f", float(msg.get("position", 0)))

            elif msg_type == "seek":
                position = float(msg.get("position", 0.0))
                logger.info("[ws] seek to position=%.3f", position)
                last_beat_reported = None  # reset so beats near seek target fire

            elif msg_type == "load":
                # Browser signals which file is now playing — update beat cache lookup
                path = msg.get("path", "")
                if path:
                    beats_cache = getattr(websocket.app.state, "beats", None) or {}
                    if path in beats_cache:
                        current_file = path
                        beat_data = beats_cache[path]
                        beats = sorted(beat_data.get("beats", []))
                        tempo = float(beat_data.get("tempo", 0.0))
                        last_beat_reported = None
                        logger.info("[ws] loaded beat data for %s: %d beats @ %.1f BPM", path, len(beats), tempo)
                    else:
                        logger.warning("[ws] no beat data cached for %s", path)

            else:
                # Unknown message type — echo back for debugging
                await websocket.send_text(json.dumps({"type": "ack", "echo": msg}))

    except WebSocketDisconnect:
        manager.disconnect(websocket)