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

from __future__ import annotations

import asyncio
import logging
from typing import Any

from yuxi.channel.context import ChannelContext
from yuxi.channel.extensions.qqbot.api_client import QQBotApiClient
from yuxi.channel.extensions.qqbot.config import QQBotConfigAdapter
from yuxi.channel.extensions.qqbot.credentials import CredentialBackup
from yuxi.channel.extensions.qqbot.session import SessionStore
from yuxi.channel.extensions.qqbot.status import QQBotStatusAdapter
from yuxi.channel.extensions.qqbot.token import TokenManager
from yuxi.channel.extensions.qqbot.types import GatewayEvent, QQBotAccountConfig
from yuxi.channel.extensions.qqbot.websocket import QQBotGatewayConnection

logger = logging.getLogger(__name__)


class QQBotGateway:
    def __init__(
        self,
        account: QQBotAccountConfig,
        config_adapter: QQBotConfigAdapter,
        status_adapter: QQBotStatusAdapter,
        message_handler: Any | None = None,
    ):
        self._account = account
        self._config_adapter = config_adapter
        self._status = status_adapter
        self._message_handler = message_handler

        self._token_manager: TokenManager | None = None
        self._api_client: QQBotApiClient | None = None
        self._session_store = SessionStore()
        self._credential_backup = CredentialBackup()
        self._connection: QQBotGatewayConnection | None = None
        self._gateway_task: asyncio.Task | None = None
        self._running = False

    async def start(self, ctx: ChannelContext) -> None:
        self._status.running = True

        if not self._account.app_id or not self._account.client_secret:
            backup = self._credential_backup.load(self._account.account_id)
            if backup:
                self._account.app_id = backup.get("app_id", "")
                self._account.client_secret = backup.get("client_secret", "")

        if not self._account.app_id or not self._account.client_secret:
            raise RuntimeError("QQBot account not configured: missing App ID or Client Secret")

        self._token_manager = TokenManager(self._account.app_id, self._account.client_secret)
        await self._token_manager.start_background_refresh()

        self._api_client = QQBotApiClient(self._token_manager)
        self._connection = QQBotGatewayConnection(
            api_client=self._api_client,
            session_store=self._session_store,
            account_id=self._account.account_id,
            dispatch_handler=self._handle_dispatch,
            reconnect_handler=self._handle_reconnect,
        )

        self._status.token_source = self._account.secret_source

        await self._connection.start()
        self._status.connected = True
        self._running = True

        self._credential_backup.save(
            self._account.account_id,
            self._account.app_id,
            self._account.client_secret or "",
        )

    async def stop(self, ctx: ChannelContext) -> None:
        self._running = False
        if self._connection:
            await self._connection.stop()
            self._connection = None
        if self._token_manager:
            await self._token_manager.close()
            self._token_manager = None
        if self._api_client:
            await self._api_client.close()
            self._api_client = None
        self._status.connected = False
        self._status.running = False

    async def _handle_dispatch(self, event: GatewayEvent) -> None:
        if self._message_handler:
            try:
                await self._message_handler(event, self._api_client, self._account)
            except Exception:
                logger.exception("[qqbot:%s] Message handler error", self._account.account_id)

    async def _handle_reconnect(self) -> None:
        logger.warning("[qqbot:%s] Handling reconnect request", self._account.account_id)
        raise RuntimeError("Gateway connection lost, triggering reconnect")

    @property
    def api_client(self) -> QQBotApiClient | None:
        return self._api_client

    @property
    def bot_openid(self) -> str | None:
        return self._connection.bot_openid if self._connection else None

    @property
    def connected(self) -> bool:
        return self._status.connected