core/token_manager.py

import time
import asyncio
from typing import Optional, Dict, Any

from core.config import ConfigManager
from core.noke_client import NokeClient

COOLDOWN_SECONDS = 300
MAX_CONSECUTIVE_ERRORS = 3

# 按错误类别冷却时间（秒）
COOLDOWN_BY_CATEGORY = {
    "auth": 3600,          # 认证失败，长冷却
    "rate_limit": 60,      # 限流，短冷却
    "signature": 60,       # 签名错误，短冷却
    "transient": 300,      # 瞬态错误，默认
}


class TokenManager:
    def __init__(self, config: ConfigManager):
        self.config = config
        self._clients: dict[str, NokeClient] = {}
        self._index: int = 0
        self._cooldowns: dict[str, float] = {}          # token_id -> 冷却截止时间戳
        self._error_categories: dict[str, str] = {}    # token_id -> 最近错误类别
        self._lock = asyncio.Lock()
        self._dirty = False

    @property
    def has_tokens(self) -> bool:
        return len(self.config.get("tokens", default=[])) > 0

    def _get_available_tokens(self) -> list[dict]:
        tokens = self.config.get("tokens", default=[])
        now = time.time()
        available = []
        for t in tokens:
            if not t.get("enabled", True):
                continue
            cooldown_until = self._cooldowns.get(t["id"], 0)
            if now >= cooldown_until:
                if t["id"] in self._cooldowns:
                    del self._cooldowns[t["id"]]
                self._error_categories.pop(t["id"], None)
                t["status"] = "unknown"
                t["error_count"] = 0
                available.append(t)
        return available

    def _get_client(self, token_info: dict) -> NokeClient:
        tid = token_info["id"]
        if tid not in self._clients:
            verify_tls = bool(self.config.get("noke", "verify_tls", default=True))
            self._clients[tid] = NokeClient(
                token_info["value"],
                self.config.get("noke", "base_url"),
                self.config.get("noke", "client_id"),
                verify_tls=verify_tls,
            )
        return self._clients[tid]

    async def get_next(self) -> tuple[Optional[dict], Optional[NokeClient]]:
        async with self._lock:
            available = self._get_available_tokens()
            if not available:
                return None, None
            self._index = self._index % len(available)
            token_info = available[self._index]
            self._index = (self._index + 1) % len(available)
            client = self._get_client(token_info)
            return token_info, client

    async def get_by_id(self, token_id: str) -> tuple[Optional[dict], Optional[NokeClient]]:
        """按 token_id 获取特定 token 和 client（用于会话粘性）"""
        async with self._lock:
            tokens = self.config.get("tokens", default=[])
            for t in tokens:
                if t["id"] == token_id:
                    # 检查是否可用（未冷却）
                    cooldown_until = self._cooldowns.get(token_id, 0)
                    if time.time() < cooldown_until:
                        return None, None
                    client = self._get_client(t)
                    return t, client
            return None, None

    def report_success(self, token_id: str):
        for t in self.config.get("tokens", default=[]):
            if t["id"] == token_id:
                t["total_requests"] = t.get("total_requests", 0) + 1
                t["error_count"] = 0
                t["last_used_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                t["status"] = "active"
                break
        self._dirty = True

    def report_error(self, token_id: str, category: str = "transient"):
        """上报错误，按类别决定冷却策略。
        - auth/rate_limit/signature: 立即进入冷却（冷却时间按类别）
        - transient: 累计错误次数，达到 MAX_CONSECUTIVE_ERRORS 才冷却（COOLDOWN_SECONDS）
        """
        for t in self.config.get("tokens", default=[]):
            if t["id"] == token_id:
                t["error_count"] = t.get("error_count", 0) + 1
                t["total_requests"] = t.get("total_requests", 0) + 1
                t["status"] = "error"
                self._error_categories[token_id] = category

                # 立即冷却类别
                if category in ("auth", "rate_limit", "signature"):
                    cooldown = COOLDOWN_BY_CATEGORY.get(category, COOLDOWN_SECONDS)
                    self._cooldowns[token_id] = time.time() + cooldown
                    t["status"] = "cooldown"
                    t["error_count"] = 0  # 重置计数
                else:
                    # transient: 累计到阈值才冷却
                    if t["error_count"] >= MAX_CONSECUTIVE_ERRORS:
                        self._cooldowns[token_id] = time.time() + COOLDOWN_SECONDS
                        t["status"] = "cooldown"
                        t["error_count"] = 0
                break
        self._dirty = True

    def remove_client(self, token_id: str):
        self._clients.pop(token_id, None)
        self._cooldowns.pop(token_id, None)
        self._error_categories.pop(token_id, None)

    def flush(self):
        if self._dirty:
            self.config.save()
            self._dirty = False

    def get_token_status(self, token_id: str) -> str:
        now = time.time()
        cooldown_until = self._cooldowns.get(token_id, 0)
        if now < cooldown_until:
            return "cooldown"
        for t in self.config.get("tokens", default=[]):
            if t["id"] == token_id:
                return t.get("status", "unknown")
        return "unknown"

    def snapshot(self) -> dict:
        """返回 token 池快照：总数、可用数、冷却中数量"""
        tokens = self.config.get("tokens", default=[])
        total = len(tokens)
        now = time.time()
        available = 0
        cooldown = 0
        for t in tokens:
            if not t.get("enabled", True):
                continue
            cooldown_until = self._cooldowns.get(t["id"], 0)
            if now < cooldown_until:
                cooldown += 1
            else:
                available += 1
        return {
            "total": total,
            "available": available,
            "cooldown": cooldown,
        }

    async def close_all(self):
        self.flush()
        for client in self._clients.values():
            await client.client.aclose()
        self._clients.clear()