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

from __future__ import annotations

import asyncio
import logging
import random
import time
from collections.abc import Callable
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass
class AccountConfig:
    account_id: str
    app_id: str
    app_secret: str
    label: str = ""
    weight: int = 1
    priority: int = 0
    group_ids: list[str] = field(default_factory=list)
    user_ids: list[str] = field(default_factory=list)
    cooldown_s: float = 30.0
    _fail_count: int = 0
    _last_fail: float = 0.0
    _last_used: float = 0.0

    @property
    def is_cooling_down(self) -> bool:
        if self._last_fail <= 0:
            return False
        return time.time() - self._last_fail < self.cooldown_s

    def record_success(self) -> None:
        self._fail_count = 0
        self._last_fail = 0.0
        self._last_used = time.time()

    def record_failure(self) -> None:
        self._fail_count += 1
        self._last_fail = time.time()

    def matches_chat(self, group_id: str = "", user_id: str = "") -> bool:
        if self.group_ids or self.user_ids:
            if group_id and self.group_ids and group_id not in self.group_ids:
                return False
            if user_id and self.user_ids and user_id not in self.user_ids:
                return False
        return True


@dataclass
class AccountRouteResult:
    account: AccountConfig
    account_id: str
    resolved: bool = True
    reason: str = ""


class MultiAccountManager:
    def __init__(
        self,
        accounts: list[AccountConfig] | None = None,
        default_rotation_strategy: str = "weighted_round_robin",
    ):
        self._accounts: dict[str, AccountConfig] = {}
        self._rotation_index = 0
        self._lock = asyncio.Lock()
        self._strategy = default_rotation_strategy
        self._route_fn: Callable | None = None

        if accounts:
            for acc in accounts:
                self._accounts[acc.account_id] = acc

    @classmethod
    def from_config(cls, config: dict | None) -> MultiAccountManager:
        if not config:
            return cls()

        accounts_cfg = config.get("accounts", [])
        if not accounts_cfg:
            app_id = config.get("app_id", "")
            app_secret = config.get("app_secret", "")
            if app_id and app_secret:
                acc = AccountConfig(
                    account_id="default",
                    app_id=app_id,
                    app_secret=app_secret,
                    label="Default",
                )
                return cls(accounts=[acc])
            return cls()

        accounts = []
        for ac in accounts_cfg:
            accounts.append(AccountConfig(
                account_id=ac.get("account_id", str(random.randint(1000, 9999))),
                app_id=ac.get("app_id", ""),
                app_secret=ac.get("app_secret", ""),
                label=ac.get("label", ""),
                weight=ac.get("weight", 1),
                priority=ac.get("priority", 0),
                group_ids=ac.get("group_ids", []),
                user_ids=ac.get("user_ids", []),
                cooldown_s=ac.get("cooldown_s", 30.0),
            ))

        return cls(accounts=accounts)

    @property
    def account_count(self) -> int:
        return len(self._accounts)

    def get_account(self, account_id: str) -> AccountConfig | None:
        return self._accounts.get(account_id)

    async def route(
        self,
        group_id: str = "",
        user_id: str = "",
        strategy: str | None = None,
    ) -> AccountRouteResult:
        async with self._lock:
            strategy = strategy or self._strategy

            if self._route_fn is not None:
                result = self._route_fn(self._accounts, group_id, user_id)
                if result:
                    return result

            candidates = [
                acc
                for acc in self._accounts.values()
                if acc.matches_chat(group_id, user_id) and not acc.is_cooling_down
            ]

            if not candidates:
                all_accounts = [
                    acc
                    for acc in self._accounts.values()
                    if acc.matches_chat(group_id, user_id)
                ]
                if all_accounts:
                    acc = all_accounts[0]
                    return AccountRouteResult(
                        account=acc,
                        account_id=acc.account_id,
                        reason="all cooling down, picked first",
                    )
                return AccountRouteResult(
                    account=AccountConfig(account_id="", app_id="", app_secret=""),
                    account_id="",
                    resolved=False,
                    reason="no matching accounts",
                )

            if strategy == "weighted_random":
                weights = [acc.weight for acc in candidates]
                total = sum(weights)
                if total <= 0:
                    acc = candidates[0]
                else:
                    r = random.uniform(0, total)
                    agg = 0
                    acc = candidates[0]
                    for candidate in candidates:
                        agg += candidate.weight
                        if r <= agg:
                            acc = candidate
                            break
            elif strategy == "least_used":
                acc = min(candidates, key=lambda a: a._last_used)
            elif strategy == "priority":
                candidates.sort(key=lambda a: (-a.priority, a._fail_count))
                acc = candidates[0]
            else:
                idx = self._rotation_index % len(candidates)
                acc = candidates[idx]
                self._rotation_index += 1

            acc.record_success()
            return AccountRouteResult(
                account=acc,
                account_id=acc.account_id,
            )

    def set_route_fn(self, fn: Callable | None) -> None:
        self._route_fn = fn

    async def mark_failure(self, account_id: str) -> None:
        async with self._lock:
            acc = self._accounts.get(account_id)
            if acc:
                acc.record_failure()
                logger.warning(
                    "MultiAccount: account %s failed (count=%d)",
                    account_id,
                    acc._fail_count,
                )

    async def mark_success(self, account_id: str) -> None:
        async with self._lock:
            acc = self._accounts.get(account_id)
            if acc:
                acc.record_success()

    async def all_cooling_down(self) -> bool:
        async with self._lock:
            return all(acc.is_cooling_down for acc in self._accounts.values()) if self._accounts else False

    def list_accounts(self) -> list[dict]:
        return [
            {
                "account_id": acc.account_id,
                "label": acc.label,
                "weight": acc.weight,
                "priority": acc.priority,
                "is_cooling_down": acc.is_cooling_down,
                "fail_count": acc._fail_count,
            }
            for acc in self._accounts.values()
        ]