from typing import List, Optional, Iterator
from loguru import logger
import threading
from ..config import settings


class KeyPoolManager:
    """Google API Key池管理器，用于处理限流重试"""
    
    def __init__(self):
        self._keys: List[str] = []
        self._current_index = 0
        self._lock = threading.Lock()
        self._initialize_keys()
    
    def _initialize_keys(self):
        """初始化key池"""
        # 解析key字符串，支持单个或多个key
        self._keys = [key.strip() for key in settings.google_api_keys.split(',') if key.strip()]
        if not self._keys:
            raise ValueError("GOOGLE_API_KEYS不能为空")
        logger.info(f"初始化Google API Key池，共{len(self._keys)}个key")
    
    def get_current_key(self) -> str:
        """获取当前key"""
        with self._lock:
            if not self._keys:
                raise ValueError("没有可用的Google API Key")
            return self._keys[self._current_index]
    
    def switch_to_next_key(self) -> bool:
        """切换到下一个key
        
        Returns:
            bool: 如果还有下一个key返回True，否则返回False
        """
        with self._lock:
            if len(self._keys) <= 1:
                return False
            
            self._current_index = (self._current_index + 1) % len(self._keys)
            logger.info(f"切换到下一个Google API Key，当前索引: {self._current_index}")
            return True
    
    def reset_to_first_key(self):
        """重置到第一个key"""
        with self._lock:
            self._current_index = 0
            logger.info("重置到第一个Google API Key")
    
    def get_all_keys(self) -> List[str]:
        """获取所有key（用于测试）"""
        return self._keys.copy()
    
    def has_multiple_keys(self) -> bool:
        """是否有多个key"""
        return len(self._keys) > 1


# 全局key池管理器实例
key_pool_manager = KeyPoolManager()