ForcePilot/backend/package/yuxi/channel/extensions/rocketchat/websocket.py

from __future__ import annotations

import asyncio
import json
import logging

import websockets

logger = logging.getLogger(__name__)

PING_INTERVAL_SECONDS = 30
SUBSCRIBE_TIMEOUT_SECONDS = 10


class RocketChatDDPClient:
    def __init__(
        self,
        server_url: str,
        user_id: str,
        auth_token: str,
        on_message=None,
    ):
        self.server_url = server_url.rstrip("/")
        self.user_id = user_id
        self.auth_token = auth_token
        self.ws_url = self.server_url.replace("http", "ws") + "/websocket"
        self._on_message = on_message
        self._ws: websockets.WebSocketClientProtocol | None = None
        self._session_id: str = ""
        self._msg_counter: int = 0
        self._pending_subs: dict[str, asyncio.Event] = {}
        self._pending_calls: dict[str, asyncio.Future] = {}
        self._active_subs: dict[str, str] = {}
        self._abort = asyncio.Event()

    @property
    def is_connected(self) -> bool:
        return self._ws is not None and not self._abort.is_set()

    async def connect(self) -> None:
        self._abort.clear()

        logger.info("Rocket.Chat DDP connecting to %s", self.ws_url)
        self._ws = await websockets.connect(self.ws_url, ping_interval=None)

        await self._send(
            {
                "msg": "connect",
                "version": "1",
                "support": ["1", "pre2", "pre1"],
            }
        )

        resp = await self._recv()
        if resp.get("msg") != "connected":
            raise ConnectionError(f"DDP handshake failed: {resp}")
        self._session_id = resp.get("session", "")
        logger.info("Rocket.Chat DDP connected, session=%s", self._session_id)

        result = await self._call("login", [{"resume": self.auth_token}])
        if not result:
            raise ConnectionError("DDP login failed")
        logger.info("Rocket.Chat DDP authenticated for user %s", self.user_id)

        await self._subscribe("stream-room-messages", ["__my_messages__", False])
        logger.info("Rocket.Chat DDP subscribed to stream-room-messages")

        try:
            await self._subscribe("stream-notify-user", [f"{self.user_id}/message", False])
            logger.info("Rocket.Chat DDP subscribed to stream-notify-user")
        except ConnectionError as e:
            logger.warning("Failed to subscribe stream-notify-user: %s", e)

        try:
            await self._subscribe(
                "stream-notify-room",
                [f"{self.user_id}/rooms-changed", False],
            )
            logger.info("Rocket.Chat DDP subscribed to stream-notify-room (rooms-changed)")
        except ConnectionError as e:
            logger.warning("Failed to subscribe stream-notify-room: %s", e)

    async def disconnect(self) -> None:
        self._abort.set()
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None

    async def listen(self, abort_event: asyncio.Event) -> None:
        ping_task = asyncio.create_task(self._ping_loop(abort_event))

        try:
            async for raw in self._ws:
                if abort_event.is_set():
                    break
                msg = json.loads(raw)
                await self._handle_frame(msg)
        except asyncio.CancelledError:
            pass
        except Exception as e:
            if not abort_event.is_set():
                logger.warning("Rocket.Chat DDP listen error: %s", e)
            raise
        finally:
            ping_task.cancel()
            try:
                await ping_task
            except asyncio.CancelledError:
                pass

    async def _handle_frame(self, msg: dict) -> None:
        msg_type = msg.get("msg", "")

        if msg_type == "ping":
            await self._send({"msg": "pong"})

        elif msg_type == "pong":
            pass

        elif msg_type == "ready":
            sub_id = msg.get("id", "")
            event = self._pending_subs.pop(sub_id, None)
            if event:
                event.set()

        elif msg_type == "result":
            call_id = msg.get("id", "")
            future = self._pending_calls.pop(call_id, None)
            if future and not future.done():
                result = msg.get("result")
                error = msg.get("error")
                if error:
                    future.set_exception(ConnectionError(f"DDP method error: {error}"))
                else:
                    future.set_result(result)

        elif msg_type == "changed":
            await self._handle_changed(msg)

        elif msg_type == "nosub":
            sub_id = msg.get("id", "")
            error = msg.get("error")
            logger.warning("DDP subscription %s failed: %s", sub_id, error)
            event = self._pending_subs.pop(sub_id, None)
            if event:
                event.set()

        elif msg_type == "added":
            pass

        elif msg_type == "removed":
            pass

    async def _handle_changed(self, msg: dict) -> None:
        collection = msg.get("collection", "")
        fields = msg.get("fields", {})
        args = fields.get("args", [])

        if collection == "stream-room-messages" and args:
            if self._on_message:
                for arg in args:
                    if isinstance(arg, dict):
                        try:
                            await self._on_message(arg)
                        except Exception as e:
                            logger.error("Error in DDP message handler: %s", e)

        elif collection == "stream-notify-room" and args:
            if self._on_message:
                for arg in args:
                    if isinstance(arg, dict):
                        try:
                            arg["_ddp_event"] = "notify-room"
                            await self._on_message(arg)
                        except Exception as e:
                            logger.error("Error in DDP notify-room handler: %s", e)

        elif collection == "stream-notify-user" and args:
            if self._on_message:
                for arg in args:
                    if isinstance(arg, dict):
                        try:
                            arg["_ddp_event"] = "notify-user"
                            await self._on_message(arg)
                        except Exception as e:
                            logger.error("Error in DDP notify-user handler: %s", e)

    async def _send(self, msg: dict) -> None:
        if self._ws:
            await self._ws.send(json.dumps(msg))

    async def _recv(self) -> dict:
        raw = await self._ws.recv()
        return json.loads(raw)

    async def _call(self, method: str, params: list) -> dict | None:
        msg_id = f"call-{self._msg_counter}"
        self._msg_counter += 1

        loop = asyncio.get_running_loop()
        future = loop.create_future()
        self._pending_calls[msg_id] = future

        await self._send(
            {
                "msg": "method",
                "method": method,
                "params": params,
                "id": msg_id,
            }
        )

        try:
            return await asyncio.wait_for(future, timeout=15.0)
        except TimeoutError:
            self._pending_calls.pop(msg_id, None)
            raise ConnectionError(f"DDP method {method} timed out")

    async def _subscribe(self, stream_name: str, params: list) -> str:
        sub_id = f"sub-{stream_name}-{self._msg_counter}"
        self._msg_counter += 1

        event = asyncio.Event()
        self._pending_subs[sub_id] = event

        await self._send(
            {
                "msg": "sub",
                "id": sub_id,
                "name": stream_name,
                "params": params,
            }
        )

        try:
            await asyncio.wait_for(event.wait(), timeout=SUBSCRIBE_TIMEOUT_SECONDS)
        except TimeoutError:
            self._pending_subs.pop(sub_id, None)
            raise ConnectionError(f"Subscription {stream_name} timed out")

        self._active_subs[sub_id] = stream_name
        return sub_id

    async def unsubscribe(self, sub_id: str) -> None:
        self._active_subs.pop(sub_id, None)
        await self._send({"msg": "unsub", "id": sub_id})
        logger.debug("DDP unsubscribed %s", sub_id)

    async def subscribe_room_event(self, room_id: str, event: str) -> str:
        return await self._subscribe("stream-notify-room", [f"{room_id}/{event}", False])

    async def _ping_loop(self, abort_event: asyncio.Event) -> None:
        while not abort_event.is_set():
            try:
                await asyncio.sleep(PING_INTERVAL_SECONDS)
                if not abort_event.is_set() and self._ws:
                    await self._send({"msg": "ping"})
            except asyncio.CancelledError:
                break
            except Exception:
                break