src/google_oauth_api.py

"""
Google OAuth2 认证模块
"""
import time
import jwt
import asyncio
from datetime import datetime, timezone, timedelta
from typing import Optional, Dict, Any, List
from urllib.parse import urlencode

from config import get_oauth_proxy_url, get_googleapis_proxy_url, get_resource_manager_api_url, get_service_usage_api_url
from log import log
from .httpx_client import get_async, post_async


class TokenError(Exception):
    """Token相关错误"""
    pass

class Credentials:
    """凭证类"""
    
    def __init__(self, access_token: str, refresh_token: str = None,
                 client_id: str = None, client_secret: str = None,
                 expires_at: datetime = None, project_id: str = None):
        self.access_token = access_token
        self.refresh_token = refresh_token
        self.client_id = client_id
        self.client_secret = client_secret
        self.expires_at = expires_at
        self.project_id = project_id
        
        # 反代配置将在使用时异步获取
        self.oauth_base_url = None
        self.token_endpoint = None
    
    def is_expired(self) -> bool:
        """检查token是否过期"""
        if not self.expires_at:
            return True
        
        # 提前3分钟认为过期
        buffer = timedelta(minutes=3)
        return (self.expires_at - buffer) <= datetime.now(timezone.utc)
    
    async def refresh_if_needed(self) -> bool:
        """如果需要则刷新token"""
        if not self.is_expired():
            return False
        
        if not self.refresh_token:
            raise TokenError("需要刷新令牌但未提供")
        
        await self.refresh()
        return True
    
    async def refresh(self, max_retries: int = 3, base_delay: float = 1.0):
        """刷新访问令牌，支持重试机制"""
        if not self.refresh_token:
            raise TokenError("无刷新令牌")
        
        data = {
            'client_id': self.client_id,
            'client_secret': self.client_secret,
            'refresh_token': self.refresh_token,
            'grant_type': 'refresh_token'
        }
        
        last_exception = None
        for attempt in range(max_retries + 1):
            try:
                oauth_base_url = await get_oauth_proxy_url()
                token_url = f"{oauth_base_url.rstrip('/')}/token"
                response = await post_async(
                    token_url,
                    data=data,
                    headers={'Content-Type': 'application/x-www-form-urlencoded'}
                )
                response.raise_for_status()
                
                token_data = response.json()
                self.access_token = token_data['access_token']
                
                if 'expires_in' in token_data:
                    expires_in = int(token_data['expires_in'])
                    self.expires_at = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
                
                if 'refresh_token' in token_data:
                    self.refresh_token = token_data['refresh_token']
                
                if attempt > 0:
                    log.debug(f"Token刷新成功（第{attempt + 1}次尝试），过期时间: {self.expires_at}")
                else:
                    log.debug(f"Token刷新成功，过期时间: {self.expires_at}")
                return
                
            except Exception as e:
                last_exception = e
                error_msg = str(e)
                
                # 检查是否是不可恢复的错误，如果是则不重试
                if self._is_non_retryable_error(error_msg):
                    log.error(f"Token刷新遇到不可恢复错误: {error_msg}")
                    break
                
                if attempt < max_retries:
                    # 计算退避延迟时间（指数退避）
                    delay = base_delay * (2 ** attempt)
                    log.warning(f"Token刷新失败（第{attempt + 1}次尝试）: {error_msg}，{delay}秒后重试...")
                    await asyncio.sleep(delay)
                else:
                    break
        
        # 所有重试都失败了
        error_msg = f"Token刷新失败（已重试{max_retries}次）: {str(last_exception)}"
        log.error(error_msg)
        raise TokenError(error_msg)
    
    def _is_non_retryable_error(self, error_msg: str) -> bool:
        """判断是否是不需要重试的错误"""
        non_retryable_patterns = [
            "400 Bad Request",
            "invalid_grant",
            "refresh_token_expired",
            "invalid_refresh_token", 
            "unauthorized_client",
            "access_denied",
            "401 Unauthorized"
        ]
        
        error_msg_lower = error_msg.lower()
        for pattern in non_retryable_patterns:
            if pattern.lower() in error_msg_lower:
                return True
                
        return False
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Credentials':
        """从字典创建凭证"""
        # 处理过期时间
        expires_at = None
        if 'expiry' in data and data['expiry']:
            try:
                expiry_str = data['expiry']
                if isinstance(expiry_str, str):
                    if expiry_str.endswith('Z'):
                        expires_at = datetime.fromisoformat(expiry_str.replace('Z', '+00:00'))
                    elif '+' in expiry_str:
                        expires_at = datetime.fromisoformat(expiry_str)
                    else:
                        expires_at = datetime.fromisoformat(expiry_str).replace(tzinfo=timezone.utc)
            except ValueError:
                log.warning(f"无法解析过期时间: {expiry_str}")
        
        return cls(
            access_token=data.get('token') or data.get('access_token', ''),
            refresh_token=data.get('refresh_token'),
            client_id=data.get('client_id'),
            client_secret=data.get('client_secret'),
            expires_at=expires_at,
            project_id=data.get('project_id')
        )
    
    def to_dict(self) -> Dict[str, Any]:
        """转为字典"""
        result = {
            'access_token': self.access_token,
            'refresh_token': self.refresh_token,
            'client_id': self.client_id,
            'client_secret': self.client_secret,
            'project_id': self.project_id
        }
        
        if self.expires_at:
            result['expiry'] = self.expires_at.isoformat()
        
        return result


class Flow:
    """OAuth流程类"""
    
    def __init__(self, client_id: str, client_secret: str, scopes: List[str],
                 redirect_uri: str = None):
        self.client_id = client_id
        self.client_secret = client_secret
        self.scopes = scopes
        self.redirect_uri = redirect_uri
        
        # 反代配置将在使用时异步获取
        self.oauth_base_url = None
        self.token_endpoint = None
        self.auth_endpoint = "https://accounts.google.com/o/oauth2/auth"
        
        self.credentials: Optional[Credentials] = None
    
    def get_auth_url(self, state: str = None, **kwargs) -> str:
        """生成授权URL"""
        params = {
            'client_id': self.client_id,
            'redirect_uri': self.redirect_uri,
            'scope': ' '.join(self.scopes),
            'response_type': 'code',
            'access_type': 'offline',
            'prompt': 'consent',
            'include_granted_scopes': 'true'
        }
        
        if state:
            params['state'] = state
        
        params.update(kwargs)
        return f"{self.auth_endpoint}?{urlencode(params)}"
    
    async def exchange_code(self, code: str) -> Credentials:
        """用授权码换取token"""
        data = {
            'client_id': self.client_id,
            'client_secret': self.client_secret,
            'redirect_uri': self.redirect_uri,
            'code': code,
            'grant_type': 'authorization_code'
        }
        
        try:
            oauth_base_url = await get_oauth_proxy_url()
            token_url = f"{oauth_base_url.rstrip('/')}/token"
            response = await post_async(
                token_url,
                data=data,
                headers={'Content-Type': 'application/x-www-form-urlencoded'}
            )
            response.raise_for_status()
            
            token_data = response.json()
            
            # 计算过期时间
            expires_at = None
            if 'expires_in' in token_data:
                expires_in = int(token_data['expires_in'])
                expires_at = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
            
            # 创建凭证对象
            self.credentials = Credentials(
                access_token=token_data['access_token'],
                refresh_token=token_data.get('refresh_token'),
                client_id=self.client_id,
                client_secret=self.client_secret,
                expires_at=expires_at
            )
            
            return self.credentials
            
        except Exception as e:
            error_msg = f"获取token失败: {str(e)}"
            log.error(error_msg)
            raise TokenError(error_msg)


class ServiceAccount:
    """Service Account类"""
    
    def __init__(self, email: str, private_key: str, project_id: str = None,
                 scopes: List[str] = None):
        self.email = email
        self.private_key = private_key
        self.project_id = project_id
        self.scopes = scopes or []
        
        # 反代配置将在使用时异步获取
        self.oauth_base_url = None
        self.token_endpoint = None
        
        self.access_token: Optional[str] = None
        self.expires_at: Optional[datetime] = None
    
    def is_expired(self) -> bool:
        """检查token是否过期"""
        if not self.expires_at:
            return True
        
        buffer = timedelta(minutes=3)
        return (self.expires_at - buffer) <= datetime.now(timezone.utc)
    
    def create_jwt(self) -> str:
        """创建JWT令牌"""
        now = int(time.time())
        
        payload = {
            'iss': self.email,
            'scope': ' '.join(self.scopes) if self.scopes else '',
            'aud': self.token_endpoint,
            'exp': now + 3600,
            'iat': now
        }
        
        return jwt.encode(payload, self.private_key, algorithm='RS256')
    
    async def get_access_token(self) -> str:
        """获取访问令牌"""
        if not self.is_expired() and self.access_token:
            return self.access_token
        
        assertion = self.create_jwt()
        
        data = {
            'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
            'assertion': assertion
        }
        
        try:
            oauth_base_url = await get_oauth_proxy_url()
            token_url = f"{oauth_base_url.rstrip('/')}/token"
            response = await post_async(
                token_url,
                data=data,
                headers={'Content-Type': 'application/x-www-form-urlencoded'}
            )
            response.raise_for_status()
            
            token_data = response.json()
            self.access_token = token_data['access_token']
            
            if 'expires_in' in token_data:
                expires_in = int(token_data['expires_in'])
                self.expires_at = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
            
            return self.access_token
            
        except Exception as e:
            error_msg = f"Service Account获取token失败: {str(e)}"
            log.error(error_msg)
            raise TokenError(error_msg)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any], scopes: List[str] = None) -> 'ServiceAccount':
        """从字典创建Service Account凭证"""
        return cls(
            email=data['client_email'],
            private_key=data['private_key'],
            project_id=data.get('project_id'),
            scopes=scopes
        )


# 工具函数
async def get_user_info(credentials: Credentials) -> Optional[Dict[str, Any]]:
    """获取用户信息"""
    await credentials.refresh_if_needed()
    
    try:
        googleapis_base_url = await get_googleapis_proxy_url()
        userinfo_url = f"{googleapis_base_url.rstrip('/')}/oauth2/v2/userinfo"
        response = await get_async(
            userinfo_url,
            headers={'Authorization': f'Bearer {credentials.access_token}'}
        )
        response.raise_for_status()
        return response.json()
    except Exception as e:
        log.error(f"获取用户信息失败: {e}")
        return None


async def get_user_email(credentials: Credentials) -> Optional[str]:
    """获取用户邮箱地址"""
    try:
        # 确保凭证有效
        await credentials.refresh_if_needed()
        
        # 调用Google userinfo API获取邮箱
        user_info = await get_user_info(credentials)
        if user_info:
            email = user_info.get("email")
            if email:
                log.info(f"成功获取邮箱地址: {email}")
                return email
            else:
                log.warning(f"userinfo响应中没有邮箱信息: {user_info}")
                return None
        else:
            log.warning("获取用户信息失败")
            return None
                
    except Exception as e:
        log.error(f"获取用户邮箱失败: {e}")
        return None


async def fetch_user_email_from_file(cred_data: Dict[str, Any]) -> Optional[str]:
    """从凭证数据获取用户邮箱地址（支持统一存储）"""
    try:
        # 直接从凭证数据创建凭证对象
        credentials = Credentials.from_dict(cred_data)
        if not credentials or not credentials.access_token:
            log.warning(f"无法从凭证数据创建凭证对象或获取访问令牌")
            return None
        
        # 获取邮箱
        return await get_user_email(credentials)
                
    except Exception as e:
        log.error(f"从凭证数据获取用户邮箱失败: {e}")
        return None


async def validate_token(token: str) -> Optional[Dict[str, Any]]:
    """验证访问令牌"""
    try:
        oauth_base_url = await get_oauth_proxy_url()
        tokeninfo_url = f"{oauth_base_url.rstrip('/')}/tokeninfo?access_token={token}"
        
        response = await get_async(tokeninfo_url)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        log.error(f"验证令牌失败: {e}")
        return None


async def enable_required_apis(credentials: Credentials, project_id: str) -> bool:
    """自动启用必需的API服务"""
    try:
        # 确保凭证有效
        if credentials.is_expired() and credentials.refresh_token:
            await credentials.refresh()
        
        headers = {
            "Authorization": f"Bearer {credentials.access_token}",
            "Content-Type": "application/json",
            "User-Agent": "geminicli-oauth/1.0",
        }
        
        # 需要启用的服务列表
        required_services = [
            "geminicloudassist.googleapis.com",  # Gemini Cloud Assist API
            "cloudaicompanion.googleapis.com"    # Gemini for Google Cloud API
        ]
        
        for service in required_services:
            log.info(f"正在检查并启用服务: {service}")
            
            # 检查服务是否已启用
            service_usage_base_url = await get_service_usage_api_url()
            check_url = f"{service_usage_base_url.rstrip('/')}/v1/projects/{project_id}/services/{service}"
            try:
                check_response = await get_async(check_url, headers=headers)
                if check_response.status_code == 200:
                    service_data = check_response.json()
                    if service_data.get("state") == "ENABLED":
                        log.info(f"服务 {service} 已启用")
                        continue
            except Exception as e:
                log.debug(f"检查服务状态失败，将尝试启用: {e}")
            
            # 启用服务
            enable_url = f"{service_usage_base_url.rstrip('/')}/v1/projects/{project_id}/services/{service}:enable"
            try:
                enable_response = await post_async(enable_url, headers=headers, json={})
                
                if enable_response.status_code in [200, 201]:
                    log.info(f"✅ 成功启用服务: {service}")
                elif enable_response.status_code == 400:
                    error_data = enable_response.json()
                    if "already enabled" in error_data.get("error", {}).get("message", "").lower():
                        log.info(f"✅ 服务 {service} 已经启用")
                    else:
                        log.warning(f"⚠️ 启用服务 {service} 时出现警告: {error_data}")
                else:
                    log.warning(f"⚠️ 启用服务 {service} 失败: {enable_response.status_code} - {enable_response.text}")
                    
            except Exception as e:
                log.warning(f"⚠️ 启用服务 {service} 时发生异常: {e}")
                
        return True
        
    except Exception as e:
        log.error(f"启用API服务时发生错误: {e}")
        return False


async def get_user_projects(credentials: Credentials) -> List[Dict[str, Any]]:
    """获取用户可访问的Google Cloud项目列表"""
    try:
        # 确保凭证有效
        if credentials.is_expired() and credentials.refresh_token:
            await credentials.refresh()
        
        headers = {
            "Authorization": f"Bearer {credentials.access_token}",
            "User-Agent": "geminicli-oauth/1.0",
        }
        
        # 使用Resource Manager API的正确域名和端点
        resource_manager_base_url = await get_resource_manager_api_url()
        url = f"{resource_manager_base_url.rstrip('/')}/v1/projects"
        log.info(f"正在调用API: {url}")
        response = await get_async(url, headers=headers)
        
        log.info(f"API响应状态码: {response.status_code}")
        if response.status_code != 200:
            log.error(f"API响应内容: {response.text}")
        
        if response.status_code == 200:
            data = response.json()
            projects = data.get('projects', [])
            # 只返回活跃的项目
            active_projects = [
                project for project in projects 
                if project.get('lifecycleState') == 'ACTIVE'
            ]
            log.info(f"获取到 {len(active_projects)} 个活跃项目")
            return active_projects
        else:
            log.warning(f"获取项目列表失败: {response.status_code} - {response.text}")
            return []
            
    except Exception as e:
        log.error(f"获取用户项目列表失败: {e}")
        return []




async def select_default_project(projects: List[Dict[str, Any]]) -> Optional[str]:
    """从项目列表中选择默认项目"""
    if not projects:
        return None
    
    # 策略1：查找显示名称或项目ID包含"default"的项目
    for project in projects:
        display_name = project.get('displayName', '').lower()
        project_id = project.get('projectId', '')
        if 'default' in display_name or 'default' in project_id.lower():
            log.info(f"选择默认项目: {project_id} ({project.get('displayName', project_id)})")
            return project_id
    
    # 策略2：选择第一个项目
    first_project = projects[0]
    project_id = first_project.get('projectId', '')
    log.info(f"选择第一个项目作为默认: {project_id} ({first_project.get('displayName', project_id)})")
    return project_id