"""
KeyManager: 多 DASHSCOPE_API_KEY 智能轮换管理器

支持 Round-Robin 轮询、429 自动冷却、401 自动禁用。
"""

import logging
import os
import threading
import time
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Optional

logger = logging.getLogger("qwen_proxy")


class KeyStatus(Enum):
    HEALTHY = "healthy"
    COOLDOWN = "cooldown"
    INVALID = "invalid"


@dataclass
class KeyState:
    key: str
    status: KeyStatus = KeyStatus.HEALTHY
    fail_count: int = 0
    total_requests: int = 0
    total_success: int = 0
    cooldown_until: float = 0.0
    last_used: float = 0.0


class KeyManager:
    """线程安全的 Key 轮换管理器"""

    def __init__(
        self,
        keys: list[str],
        cooldown_seconds: float = 60.0,
        max_retries: int = 3,
    ):
        if not keys:
            raise ValueError("至少需要提供一个 API Key")

        self._lock = threading.Lock()
        self._keys: list[KeyState] = [KeyState(key=k.strip()) for k in keys if k.strip()]
        self._index = 0
        self.cooldown_seconds = cooldown_seconds
        self.max_retries = max_retries

        if not self._keys:
            raise ValueError("至少需要提供一个有效的 API Key")

        logger.info(f"KeyManager 初始化完成，共 {len(self._keys)} 个 Key")

    @classmethod
    def from_file(cls, filepath: str | Path, **kwargs) -> "KeyManager":
        """从文件加载 Key（一行一个）"""
        path = Path(filepath)
        if not path.exists():
            raise FileNotFoundError(f"Key 文件不存在: {path}")

        keys = [
            line.strip()
            for line in path.read_text(encoding="utf-8").splitlines()
            if line.strip() and not line.strip().startswith("#")
        ]
        logger.info(f"从 {path} 加载了 {len(keys)} 个 Key")
        return cls(keys=keys, **kwargs)

    @classmethod
    def auto_discover(
        cls,
        keys_file: Optional[str | Path] = None,
        **kwargs,
    ) -> "KeyManager":
        """
        按优先级自动发现 Key:
        1. 显式传入的文件路径
        2. $CWD/.keys
        3. 环境变量 DASHSCOPE_API_KEYS (逗号分隔)
        """
        # 优先级 1: 显式文件路径
        if keys_file is not None:
            return cls.from_file(keys_file, **kwargs)

        # 优先级 2: $CWD/.keys
        cwd_keys = Path.cwd() / ".keys"
        if cwd_keys.exists():
            return cls.from_file(cwd_keys, **kwargs)

        # 优先级 3: 环境变量
        env_keys = os.environ.get("DASHSCOPE_API_KEYS", "")
        if env_keys:
            keys = [k.strip() for k in env_keys.split(",") if k.strip()]
            if keys:
                logger.info(f"从环境变量 DASHSCOPE_API_KEYS 加载了 {len(keys)} 个 Key")
                return cls(keys=keys, **kwargs)

        raise RuntimeError(
            "未找到 API Key。请通过以下方式之一提供:\n"
            "  1. keys_file 参数指定文件路径\n"
            "  2. 在当前目录创建 .keys 文件（一行一个 Key）\n"
            "  3. 设置环境变量 DASHSCOPE_API_KEYS（逗号分隔）"
        )

    def get_next_key(self) -> str:
        """获取下一个可用的 Key (Round-Robin)"""
        with self._lock:
            now = time.time()
            tried = 0
            total = len(self._keys)

            while tried < total:
                state = self._keys[self._index]
                self._index = (self._index + 1) % total

                # 跳过无效 Key
                if state.status == KeyStatus.INVALID:
                    tried += 1
                    continue

                # 冷却中的 Key，检查是否已恢复
                if state.status == KeyStatus.COOLDOWN:
                    if now >= state.cooldown_until:
                        state.status = KeyStatus.HEALTHY
                        state.fail_count = 0
                        logger.info(f"Key ...{state.key[-6:]} 冷却结束，已恢复")
                    else:
                        remaining = state.cooldown_until - now
                        logger.debug(f"Key ...{state.key[-6:]} 冷却中，剩余 {remaining:.0f}s")
                        tried += 1
                        continue

                state.total_requests += 1
                state.last_used = now
                return state.key

            raise RuntimeError("所有 API Key 均不可用（冷却中或已失效）")

    def report_success(self, key: str) -> None:
        """报告请求成功"""
        with self._lock:
            state = self._find_key(key)
            if state:
                state.fail_count = 0
                state.total_success += 1

    def report_rate_limited(self, key: str) -> None:
        """报告 429 限流，冷却该 Key"""
        with self._lock:
            state = self._find_key(key)
            if state:
                state.fail_count += 1
                state.status = KeyStatus.COOLDOWN
                state.cooldown_until = time.time() + self.cooldown_seconds
                logger.warning(
                    f"Key ...{key[-6:]} 触发限流，冷却 {self.cooldown_seconds}s "
                    f"(累计失败 {state.fail_count} 次)"
                )

    def report_invalid(self, key: str) -> None:
        """报告 401 无效 Key，永久禁用"""
        with self._lock:
            state = self._find_key(key)
            if state:
                state.status = KeyStatus.INVALID
                logger.error(f"Key ...{key[-6:]} 鉴权失败，已永久禁用")

    def get_stats(self) -> list[dict]:
        """获取所有 Key 的状态统计"""
        with self._lock:
            return [
                {
                    "key": f"...{s.key[-6:]}",
                    "status": s.status.value,
                    "requests": s.total_requests,
                    "success": s.total_success,
                    "fail_count": s.fail_count,
                }
                for s in self._keys
            ]

    @property
    def available_count(self) -> int:
        """当前可用 Key 数量"""
        with self._lock:
            now = time.time()
            return sum(
                1
                for s in self._keys
                if s.status == KeyStatus.HEALTHY
                or (s.status == KeyStatus.COOLDOWN and now >= s.cooldown_until)
            )

    def _find_key(self, key: str) -> Optional[KeyState]:
        for state in self._keys:
            if state.key == key:
                return state
        return None
