src/credential_manager.py

"""
凭证管理器 - 完全基于统一存储中间层
"""
import asyncio
import time
from datetime import datetime, timezone
from typing import Dict, Any, List, Optional, Tuple
from contextlib import asynccontextmanager

from config import get_calls_per_rotation, is_mongodb_mode
from log import log
from .storage_adapter import get_storage_adapter
from .google_oauth_api import fetch_user_email_from_file, Credentials
from .task_manager import task_manager


class CredentialManager:
    """
    统一凭证管理器
    所有存储操作通过storage_adapter进行
    """
    
    def __init__(self):
        # 核心状态
        self._initialized = False
        self._storage_adapter = None
        
        # 凭证轮换相关
        self._credential_files: List[str] = []  # 存储凭证文件名列表
        self._current_credential_index = 0
        self._call_count = 0
        self._last_scan_time = 0
        
        # 当前使用的凭证信息
        self._current_credential_file: Optional[str] = None
        self._current_credential_data: Optional[Dict[str, Any]] = None
        self._current_credential_state: Dict[str, Any] = {}
        
        # 并发控制
        self._state_lock = asyncio.Lock()
        self._operation_lock = asyncio.Lock()
        
        # 工作线程控制
        self._shutdown_event = asyncio.Event()
        self._write_worker_running = False
        self._write_worker_task = None
        
        # 原子操作计数器
        self._atomic_counter = 0
        self._atomic_lock = asyncio.Lock()
        
        # Onboarding state
        self._onboarding_complete = False
        self._onboarding_checked = False
    
    async def initialize(self):
        """初始化凭证管理器"""
        async with self._state_lock:
            if self._initialized:
                return
            
            # 初始化统一存储适配器
            self._storage_adapter = await get_storage_adapter()
            
            # 启动后台工作线程
            await self._start_background_workers()
            
            # 发现并加载凭证
            await self._discover_credentials()
            
            self._initialized = True
            storage_type = "MongoDB" if await is_mongodb_mode() else "File"
            log.debug(f"Credential manager initialized with {storage_type} storage backend")
    
    async def close(self):
        """清理资源"""
        log.debug("Closing credential manager...")
        
        # 设置关闭标志
        self._shutdown_event.set()
        
        # 等待后台任务结束
        if self._write_worker_task:
            try:
                await asyncio.wait_for(self._write_worker_task, timeout=5.0)
            except asyncio.TimeoutError:
                log.warning("Write worker task did not finish within timeout")
                if not self._write_worker_task.done():
                    self._write_worker_task.cancel()
        
        self._initialized = False
        log.debug("Credential manager closed")
    
    async def _start_background_workers(self):
        """启动后台工作线程"""
        if not self._write_worker_running:
            self._write_worker_running = True
            self._write_worker_task = task_manager.create_task(
                self._background_worker(), 
                name="credential_background_worker"
            )
    
    async def _background_worker(self):
        """后台工作线程，处理定期任务"""
        while not self._shutdown_event.is_set():
            try:
                # 每60秒检查一次凭证更新
                await asyncio.wait_for(self._shutdown_event.wait(), timeout=60.0)
                if self._shutdown_event.is_set():
                    break
                
                # 重新发现凭证（热更新）
                await self._discover_credentials()
                
            except asyncio.TimeoutError:
                # 超时是正常的，继续下一轮
                continue
            except Exception as e:
                log.error(f"Background worker error: {e}")
                await asyncio.sleep(5)  # 错误后等待5秒再继续
    
    async def _discover_credentials(self):
        """发现和加载所有可用凭证"""
        try:
            # 从存储适配器获取所有凭证
            all_credentials = await self._storage_adapter.list_credentials()
            
            # 过滤出可用的凭证（排除被禁用的）- 批量读取状态以提升性能
            available_credentials = []
            
            # 批量获取所有凭证状态，避免多次读取状态文件
            if all_credentials:
                try:
                    all_states = await self._storage_adapter.get_all_credential_states()
                    
                    for credential_name in all_credentials:
                        normalized_name = credential_name
                        # 标准化文件名以匹配状态数据中的键
                        if hasattr(self._storage_adapter._backend, '_normalize_filename'):
                            normalized_name = self._storage_adapter._backend._normalize_filename(credential_name)
                        
                        state = all_states.get(normalized_name, {})
                        if not state.get("disabled", False):
                            available_credentials.append(credential_name)
                except Exception as e:
                    log.warning(f"Failed to batch load credential states, falling back to individual checks: {e}")
                    # 如果批量读取失败，回退到逐个检查
                    for credential_name in all_credentials:
                        try:
                            state = await self._storage_adapter.get_credential_state(credential_name)
                            if not state.get("disabled", False):
                                available_credentials.append(credential_name)
                        except Exception as e2:
                            log.warning(f"Failed to check state for credential {credential_name}: {e2}")
            
            # 更新凭证列表
            old_credentials = set(self._credential_files)
            new_credentials = set(available_credentials)
            
            if old_credentials != new_credentials:
                # 记录变化（只在非初始状态时记录）
                is_initial_load = len(old_credentials) == 0
                added = new_credentials - old_credentials
                removed = old_credentials - new_credentials
                
                self._credential_files = available_credentials
                
                # 初始加载时只记录调试信息，运行时变化才记录INFO
                if not is_initial_load:
                    if added:
                        log.info(f"发现新的可用凭证: {list(added)}")
                    if removed:
                        log.info(f"移除不可用凭证: {list(removed)}")
                else:
                    # 初始加载时只记录调试信息
                    if available_credentials:
                        log.debug(f"初始加载发现 {len(available_credentials)} 个可用凭证")
                
                # 重置当前索引如果需要
                if self._current_credential_index >= len(self._credential_files):
                    self._current_credential_index = 0
            
            if not self._credential_files:
                log.warning("No available credential files found")
            else:
                log.debug(f"Available credentials: {len(self._credential_files)} files")
                
        except Exception as e:
            log.error(f"Failed to discover credentials: {e}")
    
    async def _load_current_credential(self) -> Optional[Tuple[str, Dict[str, Any]]]:
        """加载当前选中的凭证数据，包含token过期检测和自动刷新"""
        if not self._credential_files:
            return None
        
        try:
            current_file = self._credential_files[self._current_credential_index]
            
            # 从存储适配器加载凭证数据
            credential_data = await self._storage_adapter.get_credential(current_file)
            if not credential_data:
                log.error(f"Failed to load credential data for: {current_file}")
                return None
            
            # 检查refresh_token
            if "refresh_token" not in credential_data or not credential_data["refresh_token"]:
                log.warning(f"No refresh token in {current_file}")
                return None
                
            # Auto-add 'type' field if missing but has required OAuth fields
            if 'type' not in credential_data and all(key in credential_data for key in ['client_id', 'refresh_token']):
                credential_data['type'] = 'authorized_user'
                log.debug(f"Auto-added 'type' field to credential from file {current_file}")
            
            # 兼容不同的token字段格式
            if "access_token" in credential_data and "token" not in credential_data:
                credential_data["token"] = credential_data["access_token"]
            if "scope" in credential_data and "scopes" not in credential_data:
                credential_data["scopes"] = credential_data["scope"].split()
            
            # token过期检测和刷新
            should_refresh = await self._should_refresh_token(credential_data)
            
            if should_refresh:
                log.debug(f"Token需要刷新 - 文件: {current_file}")
                refreshed_data = await self._refresh_token(credential_data, current_file)
                if refreshed_data:
                    credential_data = refreshed_data
                    log.debug(f"Token刷新成功: {current_file}")
                else:
                    log.error(f"Token刷新失败: {current_file}")
                    return None
            
            # 加载状态信息
            state_data = await self._storage_adapter.get_credential_state(current_file)
            
            # 缓存当前凭证信息
            self._current_credential_file = current_file
            self._current_credential_data = credential_data
            self._current_credential_state = state_data
            
            return current_file, credential_data
            
        except Exception as e:
            log.error(f"Error loading current credential: {e}")
            return None
    
    async def get_valid_credential(self) -> Optional[Tuple[str, Dict[str, Any]]]:
        """获取有效的凭证，自动处理轮换和失效凭证切换"""
        async with self._operation_lock:
            if not self._credential_files:
                await self._discover_credentials()
                if not self._credential_files:
                    return None
            
            # 检查是否需要轮换
            if await self._should_rotate():
                await self._rotate_credential()
            
            # 尝试获取有效凭证，如果失败则自动切换
            max_attempts = len(self._credential_files)  # 最多尝试所有凭证
            
            for attempt in range(max_attempts):
                try:
                    # 加载当前凭证
                    result = await self._load_current_credential()
                    if result:
                        return result
                    
                    # 当前凭证加载失败，标记为失效并切换到下一个
                    current_file = self._credential_files[self._current_credential_index] if self._credential_files else None
                    if current_file:
                        log.warning(f"凭证失效，自动禁用并切换: {current_file}")
                        await self.set_cred_disabled(current_file, True)
                        
                        # 重新发现可用凭证（排除刚禁用的）
                        await self._discover_credentials()
                        if not self._credential_files:
                            log.error("没有可用的凭证")
                            return None
                        
                        # 重置索引到第一个可用凭证
                        self._current_credential_index = 0
                        log.info(f"切换到下一个可用凭证 (索引: {self._current_credential_index})")
                    else:
                        log.error("无法获取当前凭证文件名")
                        break
                        
                except Exception as e:
                    log.error(f"获取凭证时发生异常 (尝试 {attempt + 1}/{max_attempts}): {e}")
                    if attempt < max_attempts - 1:
                        # 切换到下一个凭证继续尝试
                        await self._rotate_credential()
                    continue
            
            log.error(f"所有 {max_attempts} 个凭证都尝试失败")
            return None
    
    async def _should_rotate(self) -> bool:
        """检查是否需要轮换凭证"""
        if not self._credential_files or len(self._credential_files) <= 1:
            return False
        
        current_calls_per_rotation = await get_calls_per_rotation()
        return self._call_count >= current_calls_per_rotation
    
    async def _rotate_credential(self):
        """轮换到下一个凭证"""
        if len(self._credential_files) <= 1:
            return
        
        self._current_credential_index = (self._current_credential_index + 1) % len(self._credential_files)
        self._call_count = 0
        
        log.info(f"Rotated to credential index {self._current_credential_index}")
    
    async def force_rotate_credential(self):
        """强制轮换到下一个凭证（用于429错误处理）"""
        async with self._operation_lock:
            if len(self._credential_files) <= 1:
                log.warning("Only one credential available, cannot rotate")
                return
            
            await self._rotate_credential()
            log.info("Forced credential rotation due to rate limit")
    
    def increment_call_count(self):
        """增加调用计数"""
        self._call_count += 1
    
    async def update_credential_state(self, credential_name: str, state_updates: Dict[str, Any]):
        """更新凭证状态"""
        try:
            # 直接通过存储适配器更新状态
            success = await self._storage_adapter.update_credential_state(credential_name, state_updates)
            
            # 如果是当前使用的凭证，更新缓存
            if credential_name == self._current_credential_file:
                self._current_credential_state.update(state_updates)
            
            if success:
                log.debug(f"Updated credential state: {credential_name}")
            else:
                log.warning(f"Failed to update credential state: {credential_name}")
                
            return success
            
        except Exception as e:
            log.error(f"Error updating credential state {credential_name}: {e}")
            return False
    
    async def set_cred_disabled(self, credential_name: str, disabled: bool):
        """设置凭证的启用/禁用状态"""
        try:
            state_updates = {"disabled": disabled}
            success = await self.update_credential_state(credential_name, state_updates)
            
            if success:
                # 如果禁用了当前正在使用的凭证，需要重新发现可用凭证
                if disabled and credential_name == self._current_credential_file:
                    await self._discover_credentials()
                    if self._credential_files:
                        await self._rotate_credential()
                
                action = "disabled" if disabled else "enabled"
                log.info(f"Credential {action}: {credential_name}")
            
            return success
            
        except Exception as e:
            log.error(f"Error setting credential disabled state {credential_name}: {e}")
            return False
    
    async def get_creds_status(self) -> Dict[str, Dict[str, Any]]:
        """获取所有凭证的状态"""
        try:
            # 从存储适配器获取所有状态
            all_states = await self._storage_adapter.get_all_credential_states()
            return all_states
            
        except Exception as e:
            log.error(f"Error getting credential statuses: {e}")
            return {}
    
    async def get_or_fetch_user_email(self, credential_name: str) -> Optional[str]:
        """获取或获取用户邮箱地址"""
        try:
            # 首先检查缓存的状态
            state = await self._storage_adapter.get_credential_state(credential_name)
            cached_email = state.get("user_email")
            
            if cached_email:
                return cached_email
            
            # 如果没有缓存，从凭证数据获取
            credential_data = await self._storage_adapter.get_credential(credential_name)
            if not credential_data:
                return None
            
            # 尝试获取邮箱
            email = await fetch_user_email_from_file(credential_data)
            
            if email:
                # 缓存邮箱地址
                await self.update_credential_state(credential_name, {"user_email": email})
                return email
            
            return None
            
        except Exception as e:
            log.error(f"Error fetching user email for {credential_name}: {e}")
            return None
    
    async def record_api_call_result(self, credential_name: str, success: bool, error_code: Optional[int] = None):
        """记录API调用结果"""
        try:
            state_updates = {}
            
            if success:
                state_updates["last_success"] = time.time()
                # 清除错误码（如果之前有的话）
                state_updates["error_codes"] = []
            elif error_code:
                # 记录错误码
                current_state = await self._storage_adapter.get_credential_state(credential_name)
                error_codes = current_state.get("error_codes", [])
                
                if error_code not in error_codes:
                    error_codes.append(error_code)
                    # 限制错误码列表长度
                    if len(error_codes) > 10:
                        error_codes = error_codes[-10:]
                    
                state_updates["error_codes"] = error_codes
            
            if state_updates:
                await self.update_credential_state(credential_name, state_updates)
                
        except Exception as e:
            log.error(f"Error recording API call result for {credential_name}: {e}")
    
    # 原子操作支持
    @asynccontextmanager
    async def _atomic_operation(self, operation_name: str):
        """原子操作上下文管理器"""
        async with self._atomic_lock:
            self._atomic_counter += 1
            operation_id = self._atomic_counter
            log.debug(f"开始原子操作[{operation_id}]: {operation_name}")
            
            try:
                yield operation_id
                log.debug(f"完成原子操作[{operation_id}]: {operation_name}")
            except Exception as e:
                log.error(f"原子操作[{operation_id}]失败: {operation_name} - {e}")
                raise
    
    async def _should_refresh_token(self, credential_data: Dict[str, Any]) -> bool:
        """检查token是否需要刷新"""
        try:
            # 如果没有access_token或过期时间，需要刷新
            if not credential_data.get("access_token") and not credential_data.get("token"):
                log.debug("没有access_token，需要刷新")
                return True
                
            expiry_str = credential_data.get("expiry")
            if not expiry_str:
                log.debug("没有过期时间，需要刷新")
                return True
                
            # 解析过期时间
            try:
                if isinstance(expiry_str, str):
                    if "+" in expiry_str:
                        file_expiry = datetime.fromisoformat(expiry_str)
                    elif expiry_str.endswith("Z"):
                        file_expiry = datetime.fromisoformat(expiry_str.replace('Z', '+00:00'))
                    else:
                        file_expiry = datetime.fromisoformat(expiry_str)
                else:
                    log.debug("过期时间格式无效，需要刷新")
                    return True
                    
                # 确保时区信息
                if file_expiry.tzinfo is None:
                    file_expiry = file_expiry.replace(tzinfo=timezone.utc)
                    
                # 检查是否还有至少5分钟有效期
                now = datetime.now(timezone.utc)
                time_left = (file_expiry - now).total_seconds()
                
                log.debug(f"Token剩余时间: {int(time_left/60)}分钟")
                
                if time_left > 300:  # 5分钟缓冲
                    return False
                else:
                    log.debug(f"Token即将过期（剩余{int(time_left/60)}分钟），需要刷新")
                    return True
                    
            except Exception as e:
                log.warning(f"解析过期时间失败: {e}，需要刷新")
                return True
                
        except Exception as e:
            log.error(f"检查token过期时出错: {e}")
            return True
    
    async def _refresh_token(self, credential_data: Dict[str, Any], filename: str) -> Optional[Dict[str, Any]]:
        """刷新token并更新存储"""
        try:
            # 创建Credentials对象
            creds = Credentials.from_dict(credential_data)
            
            # 检查是否可以刷新
            if not creds.refresh_token:
                log.error(f"没有refresh_token，无法刷新: {filename}")
                return None
                
            # 刷新token
            log.debug(f"正在刷新token: {filename}")
            await creds.refresh()
            
            # 更新凭证数据
            if creds.access_token:
                credential_data["access_token"] = creds.access_token
                # 保持兼容性
                credential_data["token"] = creds.access_token
                
            if creds.expires_at:
                credential_data["expiry"] = creds.expires_at.isoformat()
                
            # 保存到存储
            await self._storage_adapter.store_credential(filename, credential_data)
            log.info(f"Token刷新成功并已保存: {filename}")
            
            return credential_data
            
        except Exception as e:
            error_msg = str(e)
            log.error(f"Token刷新失败 {filename}: {error_msg}")
            
            # 检查是否是凭证永久失效的错误
            is_permanent_failure = self._is_permanent_refresh_failure(error_msg)
            
            if is_permanent_failure:
                log.warning(f"检测到凭证永久失效: {filename}")
                # 记录失效状态，但不在这里禁用凭证，让上层调用者处理
                await self.record_api_call_result(filename, False, 400)
            
            return None
    
    def _is_permanent_refresh_failure(self, error_msg: str) -> bool:
        """判断是否是凭证永久失效的错误"""
        # 常见的永久失效错误模式
        permanent_error_patterns = [
            "400 Bad Request",
            "invalid_grant",
            "refresh_token_expired", 
            "invalid_refresh_token",
            "unauthorized_client",
            "access_denied"
        ]
        
        error_msg_lower = error_msg.lower()
        for pattern in permanent_error_patterns:
            if pattern.lower() in error_msg_lower:
                return True
                
        return False

    # 兼容性方法 - 保持与现有代码的接口兼容
    async def _update_token_in_file(self, file_path: str, new_token: str, expires_at=None):
        """更新凭证令牌（兼容性方法）"""
        try:
            credential_data = await self._storage_adapter.get_credential(file_path)
            if not credential_data:
                log.error(f"Credential not found for token update: {file_path}")
                return False
            
            # 更新令牌数据
            credential_data["token"] = new_token
            if expires_at:
                credential_data["expiry"] = expires_at.isoformat() if hasattr(expires_at, 'isoformat') else expires_at
            
            # 保存更新后的凭证
            success = await self._storage_adapter.store_credential(file_path, credential_data)
            
            if success:
                log.debug(f"Token updated for credential: {file_path}")
            else:
                log.error(f"Failed to update token for credential: {file_path}")
            
            return success
            
        except Exception as e:
            log.error(f"Error updating token for {file_path}: {e}")
            return False


# 全局实例管理（保持兼容性）
_credential_manager: Optional[CredentialManager] = None

async def get_credential_manager() -> CredentialManager:
    """获取全局凭证管理器实例"""
    global _credential_manager
    
    if _credential_manager is None:
        _credential_manager = CredentialManager()
        await _credential_manager.initialize()
    
    return _credential_manager