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

from __future__ import annotations

import asyncio
import logging

from yuxi.channel.extensions.zulip.client import ZulipAsyncClient
from yuxi.channel.extensions.zulip.config import ZulipConfigAdapter
from yuxi.channel.extensions.zulip.errors import ZulipAPIError, ZulipQueueExpiredError
from yuxi.channel.extensions.zulip.monitor import ZulipMonitor
from yuxi.channel.runtime.backoff import BackoffConfig

logger = logging.getLogger(__name__)


class ZulipGateway:
    def __init__(self):
        self._config_adapter = ZulipConfigAdapter()
        self._queue: asyncio.Queue | None = None
        self._abort_event: asyncio.Event | None = None
        self._poll_task: asyncio.Task | None = None
        self._running = False
        self._account: dict = {}
        self._monitor: ZulipMonitor | None = None

    async def start(self, ctx) -> object:
        account = await self._resolve_account(ctx)
        self._account = account
        realm_url = account.get("realm_url", "")
        bot_email = account.get("bot_email", "")
        bot_api_key = account.get("bot_api_key", "")

        if not realm_url or not bot_email or not bot_api_key:
            return {"running": False, "reason": "not-configured"}

        self._queue = asyncio.Queue(maxsize=1000)
        self._abort_event = asyncio.Event()
        self._running = True

        account_id = account.get("account_id", "default")
        self._monitor = ZulipMonitor(account_id, account)

        self._poll_task = asyncio.create_task(self._poll_loop(account, self._queue, self._abort_event))

        return {"running": True, "account_id": account_id, "queue": self._queue}

    async def stop(self, ctx) -> None:
        self._running = False
        if self._abort_event:
            self._abort_event.set()
        if self._poll_task:
            self._poll_task.cancel()
            try:
                await self._poll_task
            except asyncio.CancelledError:
                pass
        self._queue = None
        self._abort_event = None
        self._monitor = None

    async def _poll_loop(
        self,
        account: dict,
        queue: asyncio.Queue,
        abort_event: asyncio.Event,
    ) -> None:
        account_id = account.get("account_id", "default")
        client = ZulipAsyncClient(
            realm_url=account["realm_url"],
            bot_email=account["bot_email"],
            bot_api_key=account["bot_api_key"],
        )

        backoff = BackoffConfig(base_delay=1.0, max_delay=60.0, max_retries=0, jitter=True)
        attempt = 0
        queue_id: str | None = None
        last_event_id: int = -1

        try:
            while not abort_event.is_set():
                try:
                    if queue_id is None:
                        register = await client.register_queue(
                            event_types=["message", "reaction", "update_message", "delete_message", "typing", "presence"],
                        )
                        if register.get("result") != "success":
                            delay = backoff.compute_delay(attempt)
                            attempt += 1
                            await asyncio.sleep(delay)
                            continue

                        queue_id = register["queue_id"]
                        last_event_id = register["last_event_id"]
                        attempt = 0
                        logger.info("Zulip queue registered: %s for %s", queue_id, account_id)

                    events_result = await client.get_events(
                        queue_id=queue_id,
                        last_event_id=last_event_id,
                        dont_block=False,
                    )

                    if events_result.get("result") != "success":
                        code = events_result.get("code", "")
                        if code == "BAD_EVENT_QUEUE_ID":
                            logger.warning("Zulip queue expired: %s, re-registering", queue_id)
                            queue_id = None
                            last_event_id = -1
                            continue
                        raise ZulipAPIError(500, events_result)

                    for event in events_result.get("events", []):
                        if event["type"] == "message":
                            msg = event["message"]
                            if msg.get("sender_email") == account["bot_email"]:
                                continue
                            if self._monitor:
                                unified = self._monitor.convert_zulip_message_to_unified(msg)
                            else:
                                unified = None
                            if unified and isinstance(unified, dict) and unified.get("message_id"):
                                try:
                                    queue.put_nowait(unified)
                                except asyncio.QueueFull:
                                    logger.warning("Zulip queue full, dropping message")

                        elif event["type"] == "heartbeat":
                            pass

                        elif event["type"] == "reaction":
                            if self._monitor:
                                unified = self._monitor.convert_reaction_event_to_unified(event)
                                if unified:
                                    try:
                                        queue.put_nowait(unified)
                                    except asyncio.QueueFull:
                                        logger.warning("Zulip queue full, dropping reaction event")

                        elif event["type"] == "update_message":
                            if self._monitor:
                                unified = self._monitor.convert_edit_event_to_unified(event)
                                if unified:
                                    try:
                                        queue.put_nowait(unified)
                                    except asyncio.QueueFull:
                                        logger.warning("Zulip queue full, dropping edit event")

                        elif event["type"] == "delete_message":
                            unified = {
                                "type": "message_deleted",
                                "message_ids": event.get("message_ids", []),
                                "channel_id": "zulip",
                            }
                            if unified.get("message_ids"):
                                try:
                                    queue.put_nowait(unified)
                                except asyncio.QueueFull:
                                    logger.warning("Zulip queue full, dropping delete event")

                        elif event["type"] == "typing":
                            try:
                                queue.put_nowait({"type": "typing", "data": event, "channel_id": "zulip"})
                            except asyncio.QueueFull:
                                pass

                        elif event["type"] == "presence":
                            try:
                                queue.put_nowait({"type": "presence", "data": event, "channel_id": "zulip"})
                            except asyncio.QueueFull:
                                pass

                        last_event_id = max(last_event_id, event.get("id", last_event_id))

                    attempt = 0

                except ZulipQueueExpiredError:
                    queue_id = None
                    last_event_id = -1
                    continue
                except asyncio.CancelledError:
                    break
                except Exception:
                    logger.exception("Zulip poll error for %s", account_id)
                    attempt += 1
                    delay = backoff.compute_delay(attempt)
                    await asyncio.sleep(delay)

        finally:
            if queue_id:
                try:
                    await asyncio.wait_for(client.delete_queue(queue_id), timeout=5.0)
                except Exception:
                    pass
            await client.close()
            logger.info("Zulip poll loop stopped for %s", account_id)

    async def _resolve_account(self, ctx) -> dict:
        config = getattr(ctx, "config", {}) if ctx else {}
        account_id = getattr(ctx, "account_id", "default")

        self._config_adapter.list_account_ids(config)
        return await self._config_adapter.resolve_account(account_id)