planet/motion_agent/server.py

"""Async WebSocket server for local gesture events."""

from __future__ import annotations

import asyncio
from collections.abc import Iterable
from contextlib import suppress
from typing import Any

from .cameras import (
    MotionAgentCameraError,
    NullCameraInput,
    UrlCameraInput,
    UrlCameraSpec,
    UsbCameraInput,
    UsbCameraSpec,
)
from .config import MotionAgentConfig
from .events import GestureEvent, HeartbeatEvent, SkeletonEvent, StatusEvent
from .recognizer import GestureRecognizer, MediaPipeGestureRecognizer, NullGestureRecognizer
from .state import GestureStateMachine


class MotionAgentServer:
    def __init__(
        self,
        config: MotionAgentConfig,
        cameras: Iterable[Any] | None = None,
        recognizer: GestureRecognizer | None = None,
    ) -> None:
        self.config = config
        self.cameras = list(cameras) if cameras is not None else self._build_cameras(config)
        self.recognizer = recognizer if recognizer is not None else self._build_recognizer(config)
        self.state = GestureStateMachine(
            confidence_threshold=config.confidence_threshold,
            cooldown_ms=config.cooldown_ms,
            mode=self._resolve_mode(),
        )
        self.clients: set[Any] = set()
        self.last_gesture: str | None = None
        self.last_error: str | None = None
        self._stop = asyncio.Event()

    def _build_cameras(self, config: MotionAgentConfig) -> list[Any]:
        if config.dry_run:
            return [NullCameraInput()]
        if config.camera_urls:
            return [
                UrlCameraInput(UrlCameraSpec(url=url, camera_id=f"url:{index}"))
                for index, url in enumerate(config.camera_urls)
            ]
        specs = [
            UsbCameraSpec(
                index=index,
                width=config.camera_width,
                height=config.camera_height,
                fps=config.camera_fps,
            )
            for index in config.camera_indexes
        ]
        return [UsbCameraInput(spec) for spec in specs]

    def _build_recognizer(self, config: MotionAgentConfig) -> GestureRecognizer:
        if config.dry_run:
            return NullGestureRecognizer()
        return MediaPipeGestureRecognizer()

    def _resolve_mode(self) -> str:
        if self.config.mode in {"single", "dual", "auto"}:
            if self.config.mode != "auto":
                return self.config.mode
        return "dual" if len(self.cameras) >= 2 else "single"

    def open_cameras(self) -> None:
        opened = []
        try:
            for camera in self.cameras:
                camera.open()
                opened.append(camera)
        except Exception:
            for camera in opened:
                with suppress(Exception):
                    camera.close()
            raise

    def close_cameras(self) -> None:
        for camera in self.cameras:
            with suppress(Exception):
                camera.close()

    def status_event(self, connected: bool = True) -> StatusEvent:
        return StatusEvent(
            connected=connected,
            camera_count=len(self.cameras),
            active_camera_ids=tuple(camera.camera_id for camera in self.cameras),
            mode=self.state.mode,
            recognizer=self.recognizer.name,
            last_gesture=self.last_gesture,
            error=self.last_error,
        )

    async def broadcast(
        self,
        event: GestureEvent | HeartbeatEvent | SkeletonEvent | StatusEvent,
    ) -> None:
        if not self.clients:
            return
        payload = event.to_json()
        stale = []
        for websocket in self.clients:
            try:
                await websocket.send(payload)
            except Exception:
                stale.append(websocket)
        for websocket in stale:
            self.clients.discard(websocket)

    async def handler(self, websocket: Any, path: str | None = None) -> None:
        if path is not None and path != self.config.path:
            await websocket.close(code=1008, reason="Unsupported motion agent path")
            return
        self.clients.add(websocket)
        await websocket.send(self.status_event().to_json())
        try:
            async for _message in websocket:
                await websocket.send(self.status_event().to_json())
        finally:
            self.clients.discard(websocket)

    async def heartbeat_loop(self) -> None:
        interval = max(0.1, self.config.heartbeat_interval_ms / 1000)
        while not self._stop.is_set():
            await self.broadcast(HeartbeatEvent())
            await asyncio.sleep(interval)

    async def recognition_loop(self) -> None:
        min_interval = 1 / max(1, self.config.max_event_hz)
        while not self._stop.is_set():
            try:
                for camera in self.cameras:
                    frame = camera.read()
                    observation = self.recognizer.recognize(frame)
                    matched_gesture = None
                    matched_confidence = 0.0
                    if observation is None:
                        event = None
                    else:
                        event = self.state.accept(observation)
                        if event is not None:
                            self.last_gesture = event.gesture
                            matched_gesture = event.gesture
                            matched_confidence = event.confidence
                            await self.broadcast(event)
                    skeleton = self.recognizer.debug_skeleton(
                        frame,
                        camera_id=camera.camera_id,
                        mode=self.state.mode,
                        matched_gesture=matched_gesture,
                        confidence=matched_confidence,
                    )
                    if skeleton is not None:
                        await self.broadcast(skeleton)
            except MotionAgentCameraError as exc:
                self.last_error = str(exc)
                await self.broadcast(self.status_event(connected=True))
                await asyncio.sleep(1)
            await asyncio.sleep(min_interval)

    async def run(self) -> None:
        try:
            import websockets
        except ImportError as exc:
            raise RuntimeError("websockets is required to run the motion agent server.") from exc

        self.open_cameras()
        async with websockets.serve(self.handler, self.config.host, self.config.port):
            print(f"Motion agent listening on {self.config.websocket_url}", flush=True)
            heartbeat_task = asyncio.create_task(self.heartbeat_loop())
            recognition_task = asyncio.create_task(self.recognition_loop())
            try:
                await self._stop.wait()
            finally:
                heartbeat_task.cancel()
                recognition_task.cancel()
                with suppress(asyncio.CancelledError):
                    await heartbeat_task
                with suppress(asyncio.CancelledError):
                    await recognition_task
                self.close_cameras()

    def stop(self) -> None:
        self._stop.set()