retry.py

import time
import logging
import functools
import requests
from abc import ABC, abstractmethod
from typing import Callable, Any, Dict, Optional, Type, Union, TypeVar, cast

# 导入配置模块
import config

# 类型变量定义
T = TypeVar('T')

class RetryStrategy(ABC):
    """重试策略的抽象基类"""
    
    @abstractmethod
    def should_retry(self, exception: Exception, retry_count: int, max_retries: int) -> bool:
        """
        判断是否应该重试
        
        Args:
            exception: 捕获的异常
            retry_count: 当前重试次数
            max_retries: 最大重试次数
            
        Returns:
            bool: 是否应该重试
        """
        pass
    
    @abstractmethod
    def get_retry_delay(self, retry_count: int, base_delay: int) -> float:
        """
        计算重试延迟时间
        
        Args:
            retry_count: 当前重试次数
            base_delay: 基础延迟时间（秒）
            
        Returns:
            float: 重试延迟时间（秒）
        """
        pass
    
    @abstractmethod
    def log_retry_attempt(self, logger: logging.Logger, exception: Exception, 
                         retry_count: int, max_retries: int, delay: float) -> None:
        """
        记录重试尝试
        
        Args:
            logger: 日志记录器
            exception: 捕获的异常
            retry_count: 当前重试次数
            max_retries: 最大重试次数
            delay: 重试延迟时间
        """
        pass
    
    @abstractmethod
    def on_retry(self, exception: Exception, retry_count: int) -> None:
        """
        重试前的回调函数，可以执行额外操作
        
        Args:
            exception: 捕获的异常
            retry_count: 当前重试次数
        """
        pass


class ExponentialBackoffStrategy(RetryStrategy):
    """指数退避重试策略，适用于连接错误"""
    
    def should_retry(self, exception: Exception, retry_count: int, max_retries: int) -> bool:
        return (isinstance(exception, requests.exceptions.ConnectionError) and 
                retry_count < max_retries)
    
    def get_retry_delay(self, retry_count: int, base_delay: int) -> float:
        # 指数退避: base_delay * 2^(retry_count)
        return base_delay * (2 ** retry_count)
    
    def log_retry_attempt(self, logger: logging.Logger, exception: Exception,
                         retry_count: int, max_retries: int, delay: float) -> None:
        # 检查logger是否为函数对象（如client._log）
        if callable(logger) and not isinstance(logger, logging.Logger):
            # 如果是函数，直接调用它
            logger(f"连接错误，{delay:.1f}秒后重试 ({retry_count}/{max_retries}): {exception}", "WARNING")
        else:
            # 如果是Logger对象，调用warning方法
            logger.warning(f"连接错误，{delay:.1f}秒后重试 ({retry_count}/{max_retries}): {exception}")
    
    def on_retry(self, exception: Exception, retry_count: int) -> None:
        # 连接错误不需要额外操作
        pass


class LinearBackoffStrategy(RetryStrategy):
    """线性退避重试策略，适用于超时错误"""
    
    def should_retry(self, exception: Exception, retry_count: int, max_retries: int) -> bool:
        return (isinstance(exception, requests.exceptions.Timeout) and 
                retry_count < max_retries)
    
    def get_retry_delay(self, retry_count: int, base_delay: int) -> float:
        # 线性退避: base_delay * retry_count
        return base_delay * retry_count
    
    def log_retry_attempt(self, logger: logging.Logger, exception: Exception,
                         retry_count: int, max_retries: int, delay: float) -> None:
        # 检查logger是否为函数对象（如client._log）
        if callable(logger) and not isinstance(logger, logging.Logger):
            # 如果是函数，直接调用它
            logger(f"请求超时，{delay:.1f}秒后重试 ({retry_count}/{max_retries}): {exception}", "WARNING")
        else:
            # 如果是Logger对象，调用warning方法
            logger.warning(f"请求超时，{delay:.1f}秒后重试 ({retry_count}/{max_retries}): {exception}")
    
    def on_retry(self, exception: Exception, retry_count: int) -> None:
        # 超时错误不需要额外操作
        pass


class ServerErrorStrategy(RetryStrategy):
    """服务器错误重试策略，适用于5xx错误"""
    
    def should_retry(self, exception: Exception, retry_count: int, max_retries: int) -> bool:
        if not isinstance(exception, requests.exceptions.HTTPError):
            return False
        
        response = getattr(exception, 'response', None)
        if response is None:
            return False
        
        return (500 <= response.status_code < 600 and retry_count < max_retries)
    
    def get_retry_delay(self, retry_count: int, base_delay: int) -> float:
        # 线性退避: base_delay * retry_count
        return base_delay * retry_count
    
    def log_retry_attempt(self, logger: logging.Logger, exception: Exception,
                         retry_count: int, max_retries: int, delay: float) -> None:
        response = getattr(exception, 'response', None)
        status_code = response.status_code if response else 'unknown'
        # 检查logger是否为函数对象（如client._log）
        if callable(logger) and not isinstance(logger, logging.Logger):
            # 如果是函数，直接调用它
            logger(f"服务器错误 {status_code}，{delay:.1f}秒后重试 ({retry_count}/{max_retries})", "WARNING")
        else:
            # 如果是Logger对象，调用warning方法
            logger.warning(f"服务器错误 {status_code}，{delay:.1f}秒后重试 ({retry_count}/{max_retries})")
    
    def on_retry(self, exception: Exception, retry_count: int) -> None:
        # 服务器错误不需要额外操作
        pass


class RateLimitStrategy(RetryStrategy):
    """速率限制重试策略，适用于429错误，包括账号切换逻辑和延迟重试"""
    
    def __init__(self, client=None):
        """
        初始化速率限制重试策略
        
        Args:
            client: API客户端实例，用于切换账号
        """
        self.client = client
        self.consecutive_429_count = 0  # 连续429错误计数器
    
    def should_retry(self, exception: Exception, retry_count: int, max_retries: int) -> bool:
        if not isinstance(exception, requests.exceptions.HTTPError):
            return False
        
        response = getattr(exception, 'response', None)
        if response is None:
            return False
        
        is_rate_limit = response.status_code == 429
        if is_rate_limit:
            self.consecutive_429_count += 1
        else:
            self.consecutive_429_count = 0  # 重置计数器
            
        return is_rate_limit
    
    def get_retry_delay(self, retry_count: int, base_delay: int) -> float:
        # 根据用户反馈，429错误时不需要延迟，立即重试
        return 0
    
    def log_retry_attempt(self, logger: logging.Logger, exception: Exception,
                         retry_count: int, max_retries: int, delay: float) -> None:
        # 检查logger是否为函数对象（如client._log）
        message = ""
        if self.consecutive_429_count > 1:
            message = f"连续第{self.consecutive_429_count}次速率限制错误，尝试立即重试"
        else:
            message = "速率限制错误，尝试切换账号"
            
        if callable(logger) and not isinstance(logger, logging.Logger):
            # 如果是函数，直接调用它
            logger(message, "WARNING")
        else:
            # 如果是Logger对象，调用warning方法
            logger.warning(message)
    
    def on_retry(self, exception: Exception, retry_count: int) -> None:
        # 新增: 获取关联信息
        user_identifier = getattr(self.client, '_associated_user_identifier', None)
        request_ip = getattr(self.client, '_associated_request_ip', None) # request_ip 可能在某些情况下需要

        # 只有在首次429错误或账号池中有多个账号时才切换账号
        if self.consecutive_429_count == 1 or (self.consecutive_429_count > 0 and self.consecutive_429_count % 3 == 0):
            if self.client and hasattr(self.client, 'email'):
                # 记录当前账号进入冷却期
                current_email = self.client.email # 这是切换前的 email
                config.set_account_cooldown(current_email)
                
                # 获取新账号
                new_email, new_password = config.get_next_ondemand_account_details()
                if new_email:
                    # 更新客户端信息
                    self.client.email = new_email # 这是切换后的 email
                    self.client.password = new_password
                    self.client.token = ""
                    self.client.refresh_token = ""
                    self.client.session_id = ""  # 重置会话ID，确保创建新会话
                    
                    # 尝试使用新账号登录并创建会话
                    try:
                        # 获取当前请求的上下文哈希，以便在切换账号后重新登录和创建会话时使用
                        current_context_hash = getattr(self.client, '_current_request_context_hash', None)
                        
                        self.client.sign_in(context=current_context_hash)
                        if self.client.create_session(external_context=current_context_hash):
                            # 如果成功登录并创建会话，记录日志并设置标志位
                            if hasattr(self.client, '_log'):
                                self.client._log(f"成功切换到账号 {new_email} 并使用上下文哈希 '{current_context_hash}' 重新登录和创建新会话。", "INFO")
                            # 设置标志位，通知调用方下次需要发送完整历史
                            setattr(self.client, '_new_session_requires_full_history', True)
                            if hasattr(self.client, '_log'):
                                self.client._log(f"已设置 _new_session_requires_full_history = True，下次查询应发送完整历史。", "INFO")
                        else:
                            # 会话创建失败，记录错误
                            if hasattr(self.client, '_log'):
                                self.client._log(f"切换到账号 {new_email} 后，创建新会话失败。", "WARNING")
                                # 确保在这种情况下不设置需要完整历史的标志，因为会话本身就没成功
                                setattr(self.client, '_new_session_requires_full_history', False)


                        # --- 新增: 更新 client_sessions ---
                        if not user_identifier:
                            if hasattr(self.client, '_log'):
                                self.client._log("RateLimitStrategy: _associated_user_identifier not found on client. Cannot update client_sessions.", "ERROR")
                            # 即使没有 user_identifier，账号切换和会话创建也已发生，只是无法更新全局会话池
                        else:
                            old_email_in_strategy = current_email # 切换前的 email
                            new_email_in_strategy = self.client.email # 切换后的 email (即 new_email)

                            with config.config_instance.client_sessions_lock:
                                if user_identifier in config.config_instance.client_sessions:
                                    user_specific_sessions = config.config_instance.client_sessions[user_identifier]

                                    # 1. 移除旧 email 的条目 (如果存在)
                                    #    我们只移除那些 client 实例确实是当前 self.client 的条目，
                                    #    或者更简单地，如果旧 email 存在，就移除它，因为 user_identifier
                                    #    现在应该通过 new_email 使用这个（已被修改的）client 实例。
                                    if old_email_in_strategy in user_specific_sessions:
                                        # 检查 client 实例是否匹配可能不可靠，因为 client 内部状态已变。
                                        # 直接删除旧 email 的条目，因为这个 user_identifier + client 组合现在用新 email。
                                        del user_specific_sessions[old_email_in_strategy]
                                        if hasattr(self.client, '_log'):
                                            self.client._log(f"RateLimitStrategy: Removed session for old email '{old_email_in_strategy}' for user '{user_identifier}'.", "INFO")
                                    
                                    # 2. 添加/更新新 email 的条目
                                    #    确保它指向当前这个已被修改的 self.client 实例
                                    #    并重置 active_context_hash。
                                    #    IP 地址应来自 self.client._associated_request_ip 或 routes.py 中设置的值。
                                    #    由于 routes.py 在创建/分配会话时已将 IP 存入 client_sessions，
                                    #    这里我们主要关注 client 实例和 active_context_hash。
                                    #    如果 request_ip 在 self.client 中可用，则使用它，否则尝试保留已有的。
                                    ip_to_use = request_ip if request_ip else user_specific_sessions.get(new_email_in_strategy, {}).get("ip", "unknown_ip_in_retry_update")
                                    
                                    # 需要导入 datetime
                                    from datetime import datetime
                                    
                                    # 从 client 实例获取原始请求的上下文哈希
                                    # 这个哈希应该由 routes.py 在调用 send_query 之前设置到 client 实例上
                                    active_hash_for_new_session = getattr(self.client, '_current_request_context_hash', None)

                                    user_specific_sessions[new_email_in_strategy] = {
                                        "client": self.client, # 关键: 指向当前更新了 email/session_id 的 client 实例
                                        "active_context_hash": active_hash_for_new_session, # 使用来自 client 实例的哈希
                                        "last_time": datetime.now(), # 更新时间
                                        "ip": ip_to_use
                                    }
                                    log_message_hash_part = f"set to '{active_hash_for_new_session}' (from client instance's _current_request_context_hash)" if active_hash_for_new_session is not None else "set to None (_current_request_context_hash not found on client instance)"
                                    if hasattr(self.client, '_log'):
                                        self.client._log(f"RateLimitStrategy: Updated/added session for new email '{new_email_in_strategy}' for user '{user_identifier}'. active_context_hash {log_message_hash_part}.", "INFO")
                                else:
                                    if hasattr(self.client, '_log'):
                                        self.client._log(f"RateLimitStrategy: User '{user_identifier}' not found in client_sessions during update attempt.", "WARNING")
                        # --- 更新 client_sessions 结束 ---

                    except Exception as e:
                        # 登录或创建会话失败，记录错误但不抛出异常
                        # 让后续的重试机制处理
                        if hasattr(self.client, '_log'):
                            self.client._log(f"切换到账号 {new_email} 后登录或创建会话失败: {e}", "WARNING")
                            # 此处不应更新 client_sessions，因为新账号的会话未成功建立


class RetryHandler:
    """重试处理器，管理多个重试策略"""
    
    def __init__(self, client=None, logger=None):
        """
        初始化重试处理器
        
        Args:
            client: API客户端实例，用于切换账号
            logger: 日志记录器或日志函数
        """
        self.client = client
        # 如果logger是None，使用默认logger
        # 如果logger是函数或Logger对象，直接使用
        self.logger = logger or logging.getLogger(__name__)
        self.strategies = [
            ExponentialBackoffStrategy(),
            LinearBackoffStrategy(),
            ServerErrorStrategy(),
            RateLimitStrategy(client)
        ]
    
    def retry_operation(self, operation: Callable[..., T], *args, **kwargs) -> T:
        """
        使用重试策略执行操作
        
        Args:
            operation: 要执行的操作
            *args: 操作的位置参数
            **kwargs: 操作的关键字参数
            
        Returns:
            操作的结果
            
        Raises:
            Exception: 如果所有重试都失败，则抛出最后一个异常
        """
        max_retries = config.get_config_value('max_retries')
        base_delay = config.get_config_value('retry_delay')
        retry_count = 0
        last_exception = None
        
        while True:
            try:
                return operation(*args, **kwargs)
            except Exception as e:
                last_exception = e
                
                # 查找适用的重试策略
                strategy = next((s for s in self.strategies if s.should_retry(e, retry_count, max_retries)), None)
                
                if strategy:
                    retry_count += 1
                    delay = strategy.get_retry_delay(retry_count, base_delay)
                    strategy.log_retry_attempt(self.logger, e, retry_count, max_retries, delay)
                    strategy.on_retry(e, retry_count)
                    
                    if delay > 0:
                        time.sleep(delay)
                else:
                    # 没有适用的重试策略，或者已达到最大重试次数
                    raise


def with_retry(max_retries: Optional[int] = None, retry_delay: Optional[int] = None):
    """
    重试装饰器，用于装饰需要重试的方法
    
    Args:
        max_retries: 最大重试次数，如果为None则使用配置值
        retry_delay: 基础重试延迟，如果为None则使用配置值
        
    Returns:
        装饰后的函数
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(self, *args, **kwargs):
            # 获取配置值
            _max_retries = max_retries or config.get_config_value('max_retries')
            _retry_delay = retry_delay or config.get_config_value('retry_delay')
            
            # 创建重试处理器
            handler = RetryHandler(client=self, logger=getattr(self, '_log', None))
            
            # 定义要重试的操作
            def operation():
                return func(self, *args, **kwargs)
            
            # 执行操作并处理重试
            return handler.retry_operation(operation)
        
        return wrapper
    
    return decorator