"""
API Key 负载均衡模块。

提供 Round Robin 轮询、失效 KEY 检测和冷却期管理。
"""
from typing import Dict, Optional
import time
import os
from .config import load_api_keys
from httpx import TimeoutException, ConnectError


class ApiKeyManager:
    """
    管理 API KEY 轮询、失效和冷却。
    """

    def __init__(self, keys: list[str], cooldown_seconds: int = 1800):
        """
        Args:
            keys: API KEY 列表
            cooldown_seconds: 冷却期秒数(默认 30 分钟)
        """
        self.keys = keys
        self.cooldown_seconds = cooldown_seconds
        self.key_cooldown_until: Dict[str, float] = {}
        self.current_index = 0

    def get_next_key(self) -> str:
        """
        获取下一个可用的 KEY(Round Robin)。
        跳过冷却中的 KEY。

        Returns:
            str: 可用的 API KEY

        Raises:
            RuntimeError: 所有 KEY 都在冷却中且已达到重试上限
        """
        max_retries = 3
        retry_count = 0

        while retry_count < max_retries:
            now = time.time()

            # 从 current_index 开始遍历所有 KEY
            for i in range(len(self.keys)):
                key = self.keys[(self.current_index + i) % len(self.keys)]

                # 检查是否冷却中
                if key not in self.key_cooldown_until or now > self.key_cooldown_until[key]:
                    # 找到可用的 KEY，更新索引到下一个位置
                    self.current_index = (self.current_index + i + 1) % len(self.keys)
                    return key

            # 所有 KEY 都在冷却中
            # 如果是第一次检查(current_index 已遍历一圈)或已重置为 0，需要等待
            if self.current_index == 0:
                # 等待 5 秒后重试
                if retry_count < max_retries - 1:
                    time.sleep(5)
                    retry_count += 1
                else:
                    raise RuntimeError("All API keys are in cooldown. Retry later.")
            else:
                # 重置索引并重试
                self.current_index = 0
                retry_count += 1

        raise RuntimeError("All API keys are in cooldown. Retry later.")

    def mark_key_failed(self, key: str) -> None:
        """
        标记 KEY 失败，进入冷却期。
        """
        self.key_cooldown_until[key] = time.time() + self.cooldown_seconds
        print(f"API key marked as failed. Entering cooldown for {self.cooldown_seconds}s")

    def mark_key_success(self, key: str) -> None:
        """
        标记 KEY 成功，清除冷却状态(如存在)。
        """
        if key in self.key_cooldown_until:
            del self.key_cooldown_until[key]


# 全局实例，供 app.py 使用
_api_key_manager: Optional[ApiKeyManager] = None


def get_manager() -> ApiKeyManager:
    """
    获取或创建全局管理器实例。
    """
    global _api_key_manager

    if _api_key_manager is None:
        keys = load_api_keys()
        cooldown_seconds = int(os.environ.get("SERPER_KEY_COOLDOWN_SECONDS", "1800"))
        _api_key_manager = ApiKeyManager(keys, cooldown_seconds)

    return _api_key_manager


def get_next_key() -> str:
    """
    获取下一个可用的 KEY(便捷函数)。
    """
    return get_manager().get_next_key()


def mark_key_failed(key: str) -> None:
    """
    标记 KEY 失败(便捷函数)。
    """
    get_manager().mark_key_failed(key)


def mark_key_success(key: str) -> None:
    """
    标记 KEY 成功(便捷函数)。
    """
    get_manager().mark_key_success(key)


def is_key_failure(status_code: int, exception: Optional[Exception] = None) -> bool:
    """
    判断错误是否为 KEY 相关失败。

    Args:
        status_code: HTTP 状态码
        exception: 捕获的异常

    Returns:
        bool: 如果是 KEY 相关失败返回 True，否则返回 False
    """
    # 检查异常类型
    if exception:
        if isinstance(exception, (TimeoutException, ConnectError)):
            return True

    # 检查状态码
    # 401/403: 认证/授权问题
    # 429: 速率限制
    if status_code in [401, 403, 429]:
        return True

    # 5xx: 服务器错误
    if 500 <= status_code < 600:
        return True

    # 其他 4xx 不视为 KEY 问题
    return False
