models.py

"""
账号数据模型定义 - 支持 Web Session Token (OAuth) 方式
"""
from pydantic import BaseModel, Field
from typing import Optional, Literal, Dict, Any
from datetime import datetime
import uuid


class OAuthToken(BaseModel):
    """OAuth Token 信息"""
    access_token: str
    refresh_token: str
    expires_in: int = 3600
    expiry_timestamp: int = 0  # Unix 时间戳
    project_id: Optional[str] = None
    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))


class Account(BaseModel):
    """API 账号模型 - Web Session Token 方式"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4())[:8])
    email: str
    token: OAuthToken
    enabled: bool = True
    created_at: datetime = Field(default_factory=datetime.now)
    
    # 统计信息
    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0
    last_used: Optional[datetime] = None
    last_error: Optional[str] = None
    
    # 冷却状态
    cooldown_until: Optional[datetime] = None
    
    def is_available(self) -> bool:
        """检查账号是否可用"""
        if not self.enabled:
            return False
        if self.cooldown_until and datetime.now() < self.cooldown_until:
            return False
        return True
    
    def is_token_expired(self) -> bool:
        """检查 token 是否过期（提前5分钟刷新）"""
        now = int(datetime.now().timestamp())
        return now >= self.token.expiry_timestamp - 300
    
    def display_name(self) -> str:
        """显示名称"""
        return self.email


class AccountStats(BaseModel):
    """账号统计汇总"""
    total_accounts: int = 0
    available_accounts: int = 0
    total_requests: int = 0
    success_rate: float = 0.0


class ServiceConfig(BaseModel):
    """服务配置"""
    api_key: str = "sk-antigravity"
    admin_username: str = "admin"
    admin_password: str = "antigravity"
    port: int = 7860
    enable_cors: bool = True
    max_retries: int = 3
    retry_delay: float = 1.0
    cooldown_duration: int = 60  # 秒


# ============ OpenAI 兼容格式 ============

class OpenAIMessage(BaseModel):
    """OpenAI 消息格式"""
    role: Literal["system", "user", "assistant"]
    content: str


class OpenAIChatRequest(BaseModel):
    """OpenAI Chat Completion 请求"""
    model: str
    messages: list[OpenAIMessage]
    temperature: Optional[float] = 1.0
    top_p: Optional[float] = 0.95
    max_tokens: Optional[int] = 8192
    stream: bool = False


class OpenAIChatChoice(BaseModel):
    """OpenAI 响应选项"""
    index: int = 0
    message: Optional[Dict[str, Any]] = None
    delta: Optional[Dict[str, Any]] = None
    finish_reason: Optional[str] = None


class OpenAIChatResponse(BaseModel):
    """OpenAI Chat Completion 响应"""
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:8]}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
    model: str
    choices: list[OpenAIChatChoice]
    usage: Optional[Dict[str, int]] = None


# ============ 模型映射 ============

MODEL_MAPPING: Dict[str, str] = {
    # Claude 到 Gemini 映射
    "claude-sonnet-4-5-thinking": "gemini-2.5-pro-preview",
    "claude-sonnet-4-5": "gemini-2.5-flash-preview", 
    "claude-opus-4-5": "gemini-2.5-pro-preview",
    "claude-3-5-sonnet": "gemini-2.5-flash-preview",
    "claude-3-5-haiku": "gemini-2.5-flash-lite-preview",
    
    # Gemini 模型直通
    "gemini-3-flash": "gemini-3-flash-preview",
    "gemini-3-pro": "gemini-3-pro-preview",
    "gemini-3-pro-high": "gemini-3-pro-preview",
    "gemini-2.5-pro": "gemini-2.5-pro-preview",
    "gemini-2.5-flash": "gemini-2.5-flash-preview",
    "gemini-2.5-flash-lite": "gemini-2.5-flash-lite-preview",
}

SUPPORTED_MODELS = [
    # Gemini 3 系列
    {"id": "gemini-3-flash", "name": "Gemini 3 Flash (极速预览)"},
    {"id": "gemini-3-pro", "name": "Gemini 3 Pro (经典版)"},
    {"id": "gemini-3-pro-high", "name": "Gemini 3 Pro (最强推理)"},
    
    # Gemini 2.5 系列
    {"id": "gemini-2.5-pro", "name": "Gemini 2.5 Pro (极致推理)"},
    {"id": "gemini-2.5-flash", "name": "Gemini 2.5 Flash (极速响应)"},
    {"id": "gemini-2.5-flash-lite", "name": "Gemini 2.5 Flash Lite (轻量极速)"},
    
    # Claude 映射
    {"id": "claude-sonnet-4-5-thinking", "name": "Claude Sonnet 4.5 - 思维链"},
    {"id": "claude-sonnet-4-5", "name": "Claude Sonnet 4.5"},
    {"id": "claude-opus-4-5", "name": "Claude Opus 4.5"},
]