ForcePilot/backend/package/yuxi/channels/adapters/slack/multi_account.py

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import Any

from slack_sdk.socket_mode.aiohttp import SocketModeClient
from slack_sdk.web.async_client import AsyncWebClient

from yuxi.utils.logging_config import logger


@dataclass
class SlackAccountInfo:
    account_id: str
    team_id: str = ""
    team_name: str = ""
    bot_token: str = ""
    app_token: str = ""
    signing_secret: str = ""
    bot_user_id: str = ""
    scopes: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def has_socket_mode(self) -> bool:
        return bool(self.app_token)

    @property
    def has_http_mode(self) -> bool:
        return bool(self.signing_secret)


@dataclass
class ConnectedAccount:
    account: SlackAccountInfo
    client: AsyncWebClient
    socket_handler: SocketModeClient | None = None
    socket_task: asyncio.Task | None = None
    connected: bool = False

    async def disconnect(self) -> None:
        self.connected = False
        if self.socket_handler:
            try:
                self.socket_handler.disconnect()
            except Exception:
                pass
        if self.socket_task and not self.socket_task.done():
            self.socket_task.cancel()
            try:
                await self.socket_task
            except (asyncio.CancelledError, Exception):
                pass
        self.socket_handler = None
        self.socket_task = None


@dataclass
class MultiAccountConnections:
    connections: dict[str, ConnectedAccount] = field(default_factory=dict)

    def get_connection(self, account_id: str) -> ConnectedAccount | None:
        return self.connections.get(account_id)

    def get_client(self, account_id: str) -> AsyncWebClient | None:
        conn = self.connections.get(account_id)
        return conn.client if conn else None

    def is_connected(self, account_id: str) -> bool:
        conn = self.connections.get(account_id)
        return conn.connected if conn else False

    @property
    def active_ids(self) -> list[str]:
        return [aid for aid, c in self.connections.items() if c.connected]

    async def disconnect_all(self) -> None:
        for conn in list(self.connections.values()):
            await conn.disconnect()
        self.connections.clear()


async def connect_account(account: SlackAccountInfo) -> ConnectedAccount:
    client = AsyncWebClient(token=account.bot_token)
    connected = ConnectedAccount(account=account, client=client)

    try:
        auth = await client.auth_test()
        if not auth.get("ok"):
            logger.error(f"auth.test failed for account {account.account_id}: {auth.get('error')}")
            return connected

        connected.account.bot_user_id = auth.get("user_id", "")
        connected.account.team_id = auth.get("team_id", "")

        if account.app_token:
            socket_handler = SocketModeClient(
                app_token=account.app_token,
                web_client=client,
                auto_reconnect_enabled=True,
            )
            connected.socket_handler = socket_handler
            connected.socket_task = asyncio.create_task(socket_handler.connect_async())
            connected.connected = True
            logger.info(f"Slack account {account.account_id} connected via Socket Mode")
        elif account.signing_secret:
            connected.connected = True
            logger.info(f"Slack account {account.account_id} connected via HTTP Mode")
        else:
            connected.connected = True
            logger.info(f"Slack account {account.account_id} connected (client only)")

        return connected
    except Exception as e:
        logger.error(f"Failed to connect account {account.account_id}: {e}")
        return connected


@dataclass
class MultiAccountRegistry:
    accounts: dict[str, SlackAccountInfo] = field(default_factory=dict)
    default_account_id: str = ""
    top_level_config: dict[str, Any] = field(default_factory=dict)

    def get_account(self, account_id: str | None = None) -> SlackAccountInfo | None:
        if not account_id:
            account_id = self.default_account_id
        return self.accounts.get(account_id)

    def get_merged_config(self, account_id: str | None = None) -> dict[str, Any]:
        account = self.get_account(account_id)
        merged = dict(self.top_level_config)
        if account:
            account_config = account.metadata.get("config", {}) or {}
            merged.update(account_config)
        return merged

    def register_account(self, account: SlackAccountInfo) -> None:
        self.accounts[account.account_id] = account
        if not self.default_account_id:
            self.default_account_id = account.account_id

    @classmethod
    def from_config(cls, config: dict[str, Any] | None) -> MultiAccountRegistry:
        if not config:
            return cls()
        accounts_config = config.get("accounts", []) or []
        top_level_config = dict(config)
        top_level_config.pop("accounts", None)

        if isinstance(accounts_config, dict):
            accounts = {}
            default_id = ""
            for acc_id, acc_config in accounts_config.items():
                if default_id == "":
                    default_id = acc_id
                if isinstance(acc_config, dict):
                    accounts[acc_id] = SlackAccountInfo(
                        account_id=acc_id,
                        team_id=acc_config.get("team_id", ""),
                        team_name=acc_config.get("team_name", ""),
                        bot_token=acc_config.get("bot_token", ""),
                        app_token=acc_config.get("app_token", ""),
                        signing_secret=acc_config.get("signing_secret", ""),
                        bot_user_id=acc_config.get("bot_user_id", ""),
                        scopes=list(acc_config.get("scopes", [])),
                        metadata=dict(acc_config.get("metadata", {})),
                    )
            return cls(accounts=accounts, default_account_id=default_id, top_level_config=top_level_config)

        accounts = {}
        default_id = ""
        for acc_config in accounts_config:
            if not isinstance(acc_config, dict):
                continue
            acc_id = acc_config.get("id", "")
            if default_id == "":
                default_id = acc_id
            accounts[acc_id] = SlackAccountInfo(
                account_id=acc_id,
                team_id=acc_config.get("team_id", ""),
                team_name=acc_config.get("team_name", ""),
                bot_token=acc_config.get("bot_token", ""),
                app_token=acc_config.get("app_token", ""),
                signing_secret=acc_config.get("signing_secret", ""),
                bot_user_id=acc_config.get("bot_user_id", ""),
                scopes=list(acc_config.get("scopes", [])),
                metadata=dict(acc_config.get("metadata", {})),
            )
        return cls(accounts=accounts, default_account_id=default_id, top_level_config=top_level_config)


async def _init_account_connection(
    account: SlackAccountInfo,
) -> dict[str, Any] | None:
    try:
        from slack_sdk.web.async_client import AsyncWebClient

        client = AsyncWebClient(token=account.bot_token)
        auth = await client.auth_test()
        if auth.get("ok"):
            return {
                "connected": True,
                "team_id": auth.get("team_id", account.team_id),
                "bot_user_id": auth.get("user_id", account.bot_user_id),
            }
        return {"connected": False, "error": auth.get("error")}
    except Exception as e:
        logger.error(f"Failed to init account {account.account_id}: {e}")
        return {"connected": False, "error": str(e)}