cursornew2026/backend/app/schemas/schemas.py

from pydantic import BaseModel, EmailStr, Field, model_validator
from typing import Optional, List, Any
from datetime import datetime
from app.models.models import KeyMembershipType, AccountStatus, KeyStatus


# ========== 账号相关 ==========

class AccountBase(BaseModel):
    """账号基础信息 (用于创建/更新)"""
    email: str
    token: Optional[str] = Field(None, description="兼容旧字段: user_id::jwt")
    access_token: Optional[str] = Field(None, description="Access Token")
    refresh_token: Optional[str] = Field(None, description="Refresh Token")
    workos_session_token: Optional[str] = Field(None, description="WorkosCursorSessionToken")
    password: Optional[str] = None
    remark: Optional[str] = None

class AccountCreate(AccountBase):
    """创建账号 (兼容旧字段)"""

    @model_validator(mode='before')
    @classmethod
    def ensure_token(cls, data: Any) -> Any:
        """确保至少提供一个 Token"""
        if isinstance(data, dict):
            if not data.get('token'):
                for field in ("workos_session_token", "access_token"):
                    if data.get(field):
                        data['token'] = data[field]
                        break
        return data

class AccountUpdate(BaseModel):
    email: Optional[str] = None
    token: Optional[str] = None
    access_token: Optional[str] = None
    refresh_token: Optional[str] = None
    workos_session_token: Optional[str] = None
    password: Optional[str] = None
    status: Optional[AccountStatus] = None
    remark: Optional[str] = None

class AccountResponse(BaseModel):
    """账号响应 (匹配 CursorAccount 模型)"""
    id: int
    email: str
    token: str
    access_token: Optional[str] = None
    refresh_token: Optional[str] = None
    workos_session_token: Optional[str] = None
    password: Optional[str] = None
    status: AccountStatus
    account_type: Optional[str] = None
    membership_type: Optional[str] = None
    trial_days_remaining: int = 0
    usage_limit: int = 0
    usage_used: int = 0
    usage_remaining: int = 0
    usage_percent: float = 0
    total_requests: int = 0
    locked_by_key_id: Optional[int] = None
    last_analyzed_at: Optional[datetime] = None
    remark: Optional[str] = None
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True

class AccountImport(BaseModel):
    """批量导入账号"""
    accounts: List[AccountCreate]


# ========== 外部系统批量上传 ==========

class ExternalAccountItem(BaseModel):
    """外部系统上传的账号项"""
    email: str
    access_token: str
    refresh_token: Optional[str] = None
    workos_session_token: Optional[str] = None
    membership_type: Optional[str] = "free"  # Cursor账号类型: free/free_trial/pro/business
    remark: Optional[str] = None

class ExternalBatchUpload(BaseModel):
    """外部系统批量上传请求"""
    accounts: List[ExternalAccountItem]
    update_existing: bool = True  # 是否更新已存在的账号

class ExternalBatchResponse(BaseModel):
    """外部系统批量上传响应"""
    success: bool
    total: int
    created: int
    updated: int
    failed: int
    errors: List[str] = []


# ========== 激活码相关 ==========

class KeyBase(BaseModel):
    membership_type: KeyMembershipType = KeyMembershipType.PRO  # pro=高级模型, auto=无限换号
    quota: int = 500  # 总额度 (仅Pro有效)
    valid_days: int = 30  # 有效天数，0表示永久 (仅Auto有效)
    max_devices: int = 2  # 最大设备数
    remark: Optional[str] = None

class KeyCreate(KeyBase):
    key: Optional[str] = None  # 不提供则自动生成
    count: int = 1  # 批量生成数量

class KeyUpdate(BaseModel):
    membership_type: Optional[KeyMembershipType] = None
    quota: Optional[int] = None
    valid_days: Optional[int] = None
    max_devices: Optional[int] = None
    status: Optional[KeyStatus] = None
    remark: Optional[str] = None

class KeyAddQuota(BaseModel):
    """叠加额度（只加额度不加时间）"""
    add_quota: int

class KeyResponse(BaseModel):
    id: int
    key: str
    status: KeyStatus
    membership_type: KeyMembershipType
    quota: int
    quota_used: int
    quota_remaining: Optional[int] = None  # 剩余额度（计算字段）
    valid_days: int = 30  # 有效天数 (映射自 duration_days)
    first_activated_at: Optional[datetime] = None
    expire_at: Optional[datetime] = None
    max_devices: int
    switch_count: int = 0
    last_switch_at: Optional[datetime] = None
    current_account_id: Optional[int] = None
    remark: Optional[str] = None
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True


# ========== 客户端 API 相关 ==========

class VerifyKeyRequest(BaseModel):
    key: str
    device_id: Optional[str] = None  # 设备标识

class VerifyKeyResponse(BaseModel):
    success: bool
    message: Optional[str] = None
    error: Optional[str] = None
    # 成功时返回
    email: Optional[str] = None
    accessToken: Optional[str] = None
    refreshToken: Optional[str] = None
    WorkosCursorSessionToken: Optional[str] = None
    membership_type: Optional[str] = None
    # 额度信息
    quota: Optional[int] = None  # 总额度
    quotaUsed: Optional[int] = None  # 已用额度
    quotaRemaining: Optional[int] = None  # 剩余额度
    quotaCost: Optional[int] = None  # 每次换号消耗
    expireDate: Optional[str] = None

class SwitchAccountRequest(BaseModel):
    key: str
    device_id: Optional[str] = None  # 设备标识

class SwitchAccountResponse(BaseModel):
    success: bool
    message: Optional[str] = None
    error: Optional[str] = None
    email: Optional[str] = None
    accessToken: Optional[str] = None
    refreshToken: Optional[str] = None
    WorkosCursorSessionToken: Optional[str] = None
    membership_type: Optional[str] = None
    # 额度信息
    quotaUsed: Optional[int] = None
    quotaRemaining: Optional[int] = None


# ========== 统计相关 ==========

class DashboardStats(BaseModel):
    total_accounts: int
    active_accounts: int
    pro_accounts: int
    total_keys: int
    active_keys: int
    today_usage: int


# ========== 认证相关 ==========

class Token(BaseModel):
    access_token: str
    token_type: str = "bearer"

class LoginRequest(BaseModel):
    username: str
    password: str


# ========== 全局设置相关 ==========

class GlobalSettingsResponse(BaseModel):
    """全局设置响应"""
    # ===== 密钥策略 =====
    key_max_devices: int = 2                 # 主密钥最大设备数
    auto_merge_enabled: bool = True          # 是否启用同类型密钥自动合并
    # ===== 自动检测开关 =====
    auto_analyze_enabled: bool = False       # 是否启用自动账号分析
    auto_switch_enabled: bool = True         # 是否启用自动换号
    # ===== 账号分析设置 =====
    account_analyze_interval: int = 300      # 账号分析间隔(秒)
    account_analyze_batch_size: int = 10     # 每批分析账号数量
    # ===== 换号阈值 =====
    auto_switch_threshold: int = 98          # Auto池自动换号阈值(用量百分比)
    pro_switch_threshold: int = 98           # Pro池自动换号阈值(用量百分比)
    # ===== 换号限制 =====
    max_switch_per_day: int = 50             # 每日最大换号次数
    auto_daily_switches: int = 999           # Auto密钥每日换号次数限制
    auto_switch_interval: int = 0            # Auto密钥换号冷却时间(分钟), 0表示无限制
    pro_quota_per_switch: int = 1            # Pro密钥每次换号消耗积分

class GlobalSettingsUpdate(BaseModel):
    """更新全局设置"""
    # ===== 密钥策略 =====
    key_max_devices: Optional[int] = None
    auto_merge_enabled: Optional[bool] = None
    # ===== 自动检测开关 =====
    auto_analyze_enabled: Optional[bool] = None
    auto_switch_enabled: Optional[bool] = None
    # ===== 账号分析设置 =====
    account_analyze_interval: Optional[int] = None
    account_analyze_batch_size: Optional[int] = None
    # ===== 换号阈值 =====
    auto_switch_threshold: Optional[int] = None
    pro_switch_threshold: Optional[int] = None
    # ===== 换号限制 =====
    max_switch_per_day: Optional[int] = None
    auto_daily_switches: Optional[int] = None
    auto_switch_interval: Optional[int] = None
    pro_quota_per_switch: Optional[int] = None


# ========== 批量操作相关 ==========

class BatchExtendRequest(BaseModel):
    """批量延长请求"""
    key_ids: List[int]  # 激活码ID列表
    extend_days: int = 0  # 延长天数（Auto和Pro都可用）
    add_quota: int = 0    # 增加额度（仅Pro有效）

class BatchExtendResponse(BaseModel):
    """批量延长响应"""
    success: int
    failed: int
    errors: List[str] = []


# ========== 公告相关 ==========

class AnnouncementCreate(BaseModel):
    """创建公告"""
    title: str
    content: str
    type: str = "info"  # info/warning/error/success
    is_active: bool = True

class AnnouncementUpdate(BaseModel):
    """更新公告"""
    title: Optional[str] = None
    content: Optional[str] = None
    type: Optional[str] = None
    is_active: Optional[bool] = None