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

import asyncio
import logging
from pathlib import Path

from yuxi.channel.extensions.bluebubbles.catchup import run_catchup
from yuxi.channel.extensions.bluebubbles.client import BlueBubblesClient, get_or_create_client
from yuxi.channel.extensions.bluebubbles.config import BlueBubblesConfigAdapter
from yuxi.channel.extensions.bluebubbles.dedupe import InboundDedupeStore
from yuxi.channel.extensions.bluebubbles.debounce import DebounceManager
from yuxi.channel.extensions.bluebubbles.monitor import process_inbound_message, websocket_loop
from yuxi.channel.extensions.bluebubbles.outbound import mark_chat_read
from yuxi.channel.extensions.bluebubbles.probe import probe_server, fetch_server_info
from yuxi.channel.extensions.bluebubbles.status import BlueBubblesStatus
from yuxi.channel.extensions.bluebubbles.webhook import create_webhook_router

logger = logging.getLogger(__name__)

config_adapter = BlueBubblesConfigAdapter()


class BlueBubblesGateway:
    def __init__(self):
        self._tasks: dict[str, asyncio.Task] = {}
        self._clients: dict[str, BlueBubblesClient] = {}
        self._abort_events: dict[str, asyncio.Event] = {}
        self._dedupe_stores: dict[str, InboundDedupeStore] = {}
        self._debounce_managers: dict[str, DebounceManager] = {}
        self._statuses: dict[str, BlueBubblesStatus] = {}
        self._webhook_routers: dict[str, object] = {}

    async def start_account(self, config: dict, account_id: str, ctx) -> BlueBubblesStatus:
        logger.info("Starting BlueBubbles account: %s", account_id)

        account = await config_adapter.resolve_account(config, account_id)
        if not account:
            status = BlueBubblesStatus(account_id=account_id, issues=["Not configured"])
            self._statuses[account_id] = status
            return status

        status = BlueBubblesStatus(
            account_id=account_id,
            server_url=account["server_url"],
        )
        self._statuses[account_id] = status

        client = get_or_create_client(
            account["server_url"],
            account["password"],
            account_id=account_id,
            timeout_ms=account.get("send_timeout_ms", 30000),
            allow_private_network=account.get("allow_private_network", False),
        )
        self._clients[account_id] = client

        reachable = await probe_server(client)
        status.connected = reachable
        if not reachable:
            status.issues.append("Server unreachable")
            return status

        server_info = await fetch_server_info(client)
        status.private_api_enabled = server_info.private_api_enabled
        status.imessage_logged_in = server_info.imessage_logged_in
        status.os_version = server_info.os_version
        status.macos_major = server_info.macos_major

        if not server_info.imessage_logged_in:
            status.issues.append("iMessage not logged in")

        dedupe_store = InboundDedupeStore()
        self._dedupe_stores[account_id] = dedupe_store

        debounce_manager = DebounceManager(
            window_ms=2500,
            coalesce_dms=account.get("coalesce_same_sender_dms", False),
        )
        self._debounce_managers[account_id] = debounce_manager

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

        async def _on_message(msg):
            if ctx.queue:
                await ctx.queue.put(msg)

        send_read_receipts = account.get("send_read_receipts", True)
        private_api = server_info.private_api_enabled

        async def _try_send_read_receipt(chat_guid: str):
            if send_read_receipts and private_api:
                await mark_chat_read(client, chat_guid)

        task = asyncio.create_task(
            websocket_loop(
                client,
                account_id,
                abort_event,
                dedupe_store=dedupe_store,
                debounce_manager=debounce_manager,
                on_message=_on_message,
                send_read_receipt=_try_send_read_receipt if send_read_receipts else None,
            ),
            name=f"bluebubbles-ws-{account_id}",
        )
        self._tasks[account_id] = task
        status.ws_connected = True

        webhook_secret = account.get("webhook_secret", "")
        if webhook_secret:
            self._webhook_routers[account_id] = create_webhook_router(
                path=account.get("webhook_path", "/bluebubbles-webhook"),
                webhook_secret=webhook_secret,
                account_id=account_id,
                dedupe_store=dedupe_store,
                debounce_manager=debounce_manager,
                on_message=_on_message,
                send_read_receipt=_try_send_read_receipt if send_read_receipts else None,
            )
            status.webhook_configured = True

        catchup_config = account.get("catchup")
        if catchup_config and catchup_config.enabled:
            state_dir = _resolve_state_dir(ctx)
            asyncio.create_task(
                _run_catchup_background(
                    client,
                    state_dir,
                    account_id,
                    catchup_config,
                    dedupe_store,
                    debounce_manager,
                    _on_message,
                ),
                name=f"bluebubbles-catchup-{account_id}",
            )

        return status

    async def stop_account(self, account_id: str):
        logger.info("Stopping BlueBubbles account: %s", account_id)

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

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

        client = self._clients.pop(account_id, None)
        if client:
            await client.close()

        self._dedupe_stores.pop(account_id, None)
        self._debounce_managers.pop(account_id, None)
        self._webhook_routers.pop(account_id, None)
        self._statuses.pop(account_id, None)

    def get_status(self, account_id: str) -> BlueBubblesStatus | None:
        return self._statuses.get(account_id)

    def get_webhook_router(self, account_id: str):
        return self._webhook_routers.get(account_id)


def _resolve_state_dir(ctx) -> Path:
    if hasattr(ctx, "state_dir") and ctx.state_dir:
        return Path(ctx.state_dir)
    return Path("data")


async def _run_catchup_background(
    client,
    state_dir: Path,
    account_id: str,
    catchup_config,
    dedupe_store,
    debounce_manager,
    on_message,
):
    async def _process(msg_data: dict):
        await process_inbound_message(
            msg_data,
            account_id,
            dedupe_store=dedupe_store,
            debounce_manager=debounce_manager,
            on_message=on_message,
        )

    await run_catchup(
        client,
        state_dir,
        account_id,
        catchup_config,
        process_message=_process,
        dedupe_store=dedupe_store,
    )