ForcePilot/backend/package/yuxi/channels/adapters/bluebubbles/monitor.py

from __future__ import annotations

import asyncio
import json
import random
import time
from collections.abc import Callable, Coroutine
from typing import Any

from yuxi.channels.adapters.bluebubbles.exceptions import BlueBubblesConnectionError
from yuxi.utils.logging_config import logger

try:
    import websockets.asyncio.client as ws_client
except ImportError:
    ws_client = None


class InboundReorderBuffer:
    MAX_BUFFER_MS = 500

    def __init__(self, max_buffer_ms: int | None = None):
        self._buffer: list[dict[str, Any]] = []
        self._flush_task: asyncio.Task | None = None
        self._max_buffer_ms = max_buffer_ms or self.MAX_BUFFER_MS
        self._on_flush: Callable[[list[dict[str, Any]]], Coroutine[Any, Any, None]] | None = None

    def set_on_flush(self, handler: Callable[[list[dict[str, Any]]], Coroutine[Any, Any, None]]) -> None:
        self._on_flush = handler

    def enqueue(self, event: dict[str, Any]) -> None:
        self._buffer.append(event)
        if self._flush_task is None or self._flush_task.done():
            self._flush_task = asyncio.create_task(self._delayed_flush())

    async def _delayed_flush(self) -> None:
        await asyncio.sleep(self._max_buffer_ms / 1000.0)
        await self.flush()

    async def flush(self) -> None:
        if not self._buffer:
            return
        events = self._buffer[:]
        self._buffer.clear()
        events.sort(key=lambda e: self._extract_timestamp(e))
        if self._on_flush:
            for event in events:
                await self._on_flush(event)

    @staticmethod
    def _extract_timestamp(event: dict[str, Any]) -> float:
        data = event.get("data", {})
        message = data.get("message", {})
        ts = message.get("dateCreated", 0)
        if isinstance(ts, (int, float)) and ts > 0:
            return float(ts)
        return time.time() * 1000


class BlueBubblesMonitor:
    def __init__(
        self,
        server_url: str,
        password: str,
        on_event: Callable[[dict[str, Any]], Coroutine[Any, Any, None]],
        on_disconnect: Callable[[], Coroutine[Any, Any, None]] | None = None,
        on_reconnect: Callable[[], Coroutine[Any, Any, None]] | None = None,
    ):
        self._password = password
        self._on_event = on_event
        self._on_disconnect = on_disconnect
        self._on_reconnect = on_reconnect
        ws_base = server_url.replace("http://", "ws://").replace("https://", "wss://")
        self._ws_url = f"{ws_base.rstrip('/')}/ws?password={self._password}"
        self._ws = None
        self._running = False
        self._reconnect_delay = 1.0
        self._max_reconnect_delay = 60.0
        self._task: asyncio.Task | None = None

    async def start(self) -> None:
        if ws_client is None:
            raise BlueBubblesConnectionError("websockets library not installed")
        self._running = True
        self._reconnect_delay = 1.0
        self._task = asyncio.create_task(self._listen_loop())
        logger.info("[BlueBubbles] WebSocket monitor started")

    async def stop(self) -> None:
        self._running = False
        if self._task and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        if self._ws:
            await self._ws.close()
            self._ws = None
        logger.info("[BlueBubbles] WebSocket monitor stopped")

    async def _listen_loop(self) -> None:
        while self._running:
            try:
                async for ws in ws_client.connect(self._ws_url):
                    self._ws = ws
                    self._reconnect_delay = 1.0
                    logger.info(f"[BlueBubbles] WebSocket connected: {self._ws_url}")
                    if self._on_reconnect:
                        await self._on_reconnect()
                    await self._handle_messages(ws)
            except Exception as e:
                if not self._running:
                    break
                if self._on_disconnect:
                    try:
                        await self._on_disconnect()
                    except Exception:
                        logger.exception("[BlueBubbles] on_disconnect callback failed")
                jitter = random.uniform(0, self._reconnect_delay * 0.5)
                delay = self._reconnect_delay + jitter
                logger.warning(f"[BlueBubbles] WebSocket disconnected, reconnecting in {delay:.1f}s: {e}")
                await asyncio.sleep(delay)
                self._reconnect_delay = min(
                    self._reconnect_delay * 2,
                    self._max_reconnect_delay,
                )

    async def _handle_messages(self, ws) -> None:
        async for raw in ws:
            try:
                event = json.loads(raw)
                await self._on_event(event)
            except json.JSONDecodeError:
                logger.warning("[BlueBubbles] Received invalid JSON from WebSocket")
            except Exception:
                logger.exception("[BlueBubbles] Error handling WebSocket event")