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

from __future__ import annotations

import asyncio
import logging
import random

from yuxi.channel.extensions.rocketchat.client import RocketChatClient
from yuxi.channel.extensions.rocketchat.config import (
    RocketChatConfigAdapter,
    normalize_rocketchat_server_url,
)
from yuxi.channel.extensions.rocketchat.errors import RocketChatError, RocketChatAuthError
from yuxi.channel.extensions.rocketchat.monitor import RocketChatMonitor
from yuxi.channel.extensions.rocketchat.status import RocketChatStatusAdapter
from yuxi.channel.extensions.rocketchat.websocket import RocketChatDDPClient

logger = logging.getLogger(__name__)

GATEWAY_AUTH_BYPASS_PATHS = [
    "/api/channels/rocketchat/webhook",
]


class RocketChatGatewayAdapter:
    def __init__(self, config_adapter: RocketChatConfigAdapter):
        self.config_adapter = config_adapter
        self._clients: dict[str, RocketChatClient] = {}
        self._ddp_clients: dict[str, RocketChatDDPClient] = {}
        self._monitors: dict[str, RocketChatMonitor] = {}
        self._tasks: dict[str, asyncio.Task] = {}
        self._abort_events: dict[str, asyncio.Event] = {}
        self._status_adapters: dict[str, RocketChatStatusAdapter] = {}
        self.bot_info: dict[str, dict] = {}

    async def start(self, ctx: object) -> object:
        account_id = getattr(ctx, "account_id", "default") if ctx else "default"

        account = await self.config_adapter.resolve_account(account_id)
        if not account.get("auth_token") or not account.get("user_id") or not account.get("server_url"):
            logger.warning("Rocket.Chat account %s not configured, skipping", account_id)
            return {"status": "not_configured", "account_id": account_id}

        server_url = normalize_rocketchat_server_url(account["server_url"])
        client = RocketChatClient(
            server_url=server_url,
            user_id=account["user_id"],
            auth_token=account["auth_token"],
        )

        try:
            me = await client.fetch_me()
            self.bot_info[account_id] = me
            logger.info(
                "Rocket.Chat bot %s (%s) connected to %s",
                me.get("username", "unknown"),
                me.get("_id", "unknown"),
                server_url,
            )
        except RocketChatAuthError as e:
            logger.error("Rocket.Chat auth failed for account %s: %s", account_id, e)
            await client.close()
            return {"status": "auth_failed", "account_id": account_id}
        except RocketChatError as e:
            logger.error("Rocket.Chat connection failed for account %s: %s", account_id, e)
            await client.close()
            return {"status": "connection_failed", "account_id": account_id}

        self._clients[account_id] = client
        self._status_adapters[account_id] = RocketChatStatusAdapter(client, account_id)

        monitor = RocketChatMonitor(client, account_id, account)
        monitor.bot_user_id = me.get("_id", "")
        monitor.bot_username = me.get("username", "")
        self._monitors[account_id] = monitor

        abort_event = asyncio.Event()
        self._abort_events[account_id] = abort_event

        async def ddp_connect_loop():
            ddp_client = RocketChatDDPClient(
                server_url=server_url,
                user_id=account["user_id"],
                auth_token=account["auth_token"],
                on_message=monitor.handle_message,
            )
            self._ddp_clients[account_id] = ddp_client

            delay = 2000
            max_delay = 60000
            while not abort_event.is_set():
                try:
                    await ddp_client.connect()
                    delay = 2000
                    try:
                        await client.set_user_status("online", "Bot is ready")
                        logger.info("Rocket.Chat bot status set to online for account %s", account_id)
                    except RocketChatError as e:
                        logger.warning("Failed to set bot status to online: %s", e)
                    await ddp_client.listen(abort_event)
                except asyncio.CancelledError:
                    break
                except Exception as e:
                    if abort_event.is_set():
                        break
                    logger.warning("Rocket.Chat DDP connection error: %s. Reconnecting...", e)
                    delay = min(delay * 2, max_delay)
                    spread = delay * 0.2
                    offset = random.uniform(-spread, spread)
                    jittered = max(1, round(delay + offset))
                    logger.info("Reconnecting in %d ms", jittered)
                    await asyncio.sleep(jittered / 1000)

        task = asyncio.create_task(ddp_connect_loop())
        self._tasks[account_id] = task

        return {
            "status": "started",
            "account_id": account_id,
            "server_url": server_url,
            "bot_user_id": me.get("_id", ""),
            "bot_username": me.get("username", ""),
        }

    async def stop(self, ctx: object) -> None:
        account_id = getattr(ctx, "account_id", "default") if ctx else "default"

        abort = self._abort_events.pop(account_id, None)
        if abort:
            abort.set()

        task = self._tasks.pop(account_id, None)
        if task:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

        ddp = self._ddp_clients.pop(account_id, None)
        if ddp:
            await ddp.disconnect()

        self._monitors.pop(account_id, None)

        client = self._clients.pop(account_id, None)
        if client:
            try:
                await client.set_user_status("offline", "Bot is offline")
            except Exception:
                pass
            await client.close()

        self._status_adapters.pop(account_id, None)
        self.bot_info.pop(account_id, None)
        logger.info("Rocket.Chat gateway stopped for account %s", account_id)

    def resolve_gateway_auth_bypass_paths(self, config: dict) -> list[str]:
        return GATEWAY_AUTH_BYPASS_PATHS

    def get_client(self, account_id: str = "default") -> RocketChatClient | None:
        return self._clients.get(account_id)

    def get_ddp_client(self, account_id: str = "default") -> RocketChatDDPClient | None:
        return self._ddp_clients.get(account_id)

    def get_monitor(self, account_id: str = "default") -> RocketChatMonitor | None:
        return self._monitors.get(account_id)

    def get_bot_user_id(self, account_id: str = "default") -> str:
        info = self.bot_info.get(account_id, {})
        return info.get("_id", "")

    def get_bot_username(self, account_id: str = "default") -> str:
        info = self.bot_info.get(account_id, {})
        return info.get("username", "")

    async def handle_webhook(self, request_data: bytes, headers: dict) -> dict | None:
        from yuxi.channel.extensions.rocketchat.webhook import handle_rocketchat_webhook

        account_id = "default"
        account = await self.config_adapter.resolve_account(account_id)
        webhook_secret = account.get("webhook_secret", "")
        return await handle_rocketchat_webhook(request_data, headers, webhook_secret)