taiji-AI-PAD/services/mcp-server/app/routes/auth.py

"""
认证与权限管理路由
"""

from datetime import timedelta
from fastapi import APIRouter, Depends, HTTPException, status, Query, Request
from sqlalchemy import select, and_
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.exc import IntegrityError
import secrets
import hashlib
from typing import Optional

from database import get_db
from models import (
    User, Channel, APIKey, Balance, TenantCustomAgentQuota,
    PlatformAgentQuota, TenantModelKey, ResourceAllocation, ModelProvider
)
from app.auth import (
    authenticate_user,
    create_access_token,
    create_refresh_token,
    get_password_hash,
    verify_password,
    require_auth,
)
from app.schemas import (
    LoginRequest,
    SuccessResponse,
    TokenResponse,
    PasswordChangeRequest,
    ForgotPasswordRequest,
    ResetPasswordRequest,
    APIKeyInfo,
    RegenerateAPIKeyResponse,
    UserCreate,
)
from app.email_verification import verify_code, peek_verification_code, send_and_store_verification_code, check_rate_limit, send_password_reset_code
from app.audit import log_audit_event
from pydantic import BaseModel
from app.agent_manager_client import get_agent_manager_client, AgentManagerError
from app.litellm_client import get_litellm_client, LiteLLMClientError
from config import settings
import uuid
from datetime import datetime
from sqlalchemy import and_
import structlog

logger = structlog.get_logger(__name__)

router = APIRouter(prefix="/api/auth", tags=["认证"])


class _LoginRateLimit:
    """每 IP 维度对登录请求限流（默认 5 次/60s，超出 429）。
    多副本部署下每副本独立计数，可接受。"""
    def __init__(self, max_attempts: int = 5, window_seconds: int = 60):
        self.max_attempts = max_attempts
        self.window_seconds = window_seconds
        self._attempts: dict = {}

    def _client_ip(self, request) -> str:
        xff = request.headers.get("x-forwarded-for")
        if xff:
            return xff.split(",")[0].strip()
        return request.client.host if request.client else "unknown"

    async def __call__(self, request: Request):
        import time
        from collections import deque
        ip = self._client_ip(request)
        now = time.time()
        bucket = self._attempts.setdefault(ip, deque())
        while bucket and bucket[0] < now - self.window_seconds:
            bucket.popleft()
        if len(bucket) >= self.max_attempts:
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail="登录尝试过于频繁，请稍后再试",
                headers={"Retry-After": str(self.window_seconds)},
            )
        bucket.append(now)


_login_rate_limit = _LoginRateLimit()


def _mask_api_key(key: str) -> str:
    """隐藏API密钥的中间部分"""
    if len(key) <= 12:
        return key[:4] + "..." + key[-4:]
    return key[:8] + "..." + key[-4:]


@router.post("/login", response_model=SuccessResponse)
async def login(req: LoginRequest, request: Request, db: AsyncSession = Depends(get_db), _: None = Depends(_login_rate_limit)):
    """
    用户/渠道/管理员/供应商登录

    支持的角色：
    - user: 租户用户
    - channel: 渠道管理员
    - billing_admin: 计费管理员
    - operations_admin: 运营管理员
    - admin: 管理员
    - super_admin: 超级管理员
    - provider: 供应商管理员
    """
    success = False
    result_user_id: Optional[str] = None
    error_msg: Optional[str] = None
    try:
        # 根据角色查找用户
        if req.role == "channel":
            # 渠道管理员登录
            result = await db.execute(select(Channel).where(Channel.email == req.email))
            entity = result.scalar_one_or_none()

            if not entity or not verify_password(req.password, entity.password_hash):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="邮箱或密码错误"
                )

            # 更新最后登录时间
            from datetime import datetime
            entity.last_login_at = datetime.utcnow()
            await db.commit()

            # 创建JWT token
            token_data = {
                "sub": str(entity.id),
                "email": entity.email,
                "role": "channel_admin",
                "channelId": str(entity.id),
            }
            access_token = create_access_token(data=token_data)
            refresh_token = create_refresh_token(data=token_data)

            success = True
            result_user_id = str(entity.id)
            return SuccessResponse(
                data={
                    "token": access_token,
                    "refreshToken": refresh_token,
                    "user": {
                        "id": str(entity.id),
                        "name": entity.name,
                        "email": entity.email,
                        "role": "channel_admin",
                        "channelId": str(entity.id),
                    }
                }
            )

        else:
            # 用户/管理员/供应商登录
            result = await db.execute(select(User).where(User.email == req.email))
            user = result.scalar_one_or_none()

            if not user:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="邮箱或密码错误"
                )

            # 验证密码（兼容两种密码字段）
            password_hash = user.password_hash or user.hashed_password
            if not password_hash or not verify_password(req.password, password_hash):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="邮箱或密码错误"
                )

            # 验证用户状态
            if hasattr(user, 'status') and user.status == "inactive":
                raise HTTPException(
                    status_code=status.HTTP_403_FORBIDDEN,
                    detail="账户已被禁用"
                )

            # 验证角色
            user_role = user.role

            # 角色验证逻辑
            valid_roles = {
                "super_admin": ["super_admin"],
                "admin": ["admin", "super_admin"],
                "billing_admin": ["billing_admin", "admin", "super_admin"],
                "operations_admin": ["operations_admin", "admin", "super_admin"],
                "user": ["user"],
                "provider": ["provider_admin"],
            }

            allowed_roles = valid_roles.get(req.role, [])
            if user_role not in allowed_roles:
                raise HTTPException(
                    status_code=status.HTTP_403_FORBIDDEN,
                    detail=f"权限不足，当前角色: {user_role}"
                )

            # 更新最后登录时间
            from datetime import datetime
            user.last_login_at = datetime.utcnow()
            await db.commit()

            # 创建JWT token
            token_data = {
                "sub": str(user.id),
                "email": user.email,
                "role": user_role,
                "channelId": str(user.channel_id) if user.channel_id else None,
            }
            access_token = create_access_token(data=token_data)
            refresh_token = create_refresh_token(data=token_data)

            success = True
            result_user_id = str(user.id)
            return SuccessResponse(
                data={
                    "token": access_token,
                    "refreshToken": refresh_token,
                    "user": {
                        "id": str(user.id),
                        "name": user.name or user.full_name,
                        "email": user.email,
                        "role": user_role,
                        "channelId": str(user.channel_id) if user.channel_id else None,
                    }
                }
            )
    except HTTPException as e:
        error_msg = e.detail if isinstance(e.detail, str) else str(e.detail)
        raise
    finally:
        try:
            await log_audit_event(
                action="auth.login",
                resource_type="user",
                resource_id=req.email,
                user_id=result_user_id,
                success=success,
                details={"role": req.role},
                error_message=error_msg,
                request=request,
                db=db,
            )
        except Exception as audit_exc:
            logger.warning("auth_login_audit_failed", error=str(audit_exc))


@router.get("/me", response_model=SuccessResponse)
async def get_current_user(
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db),
):
    """
    获取当前登录用户信息
    """
    user_id = principal.get("user_id")
    result = await db.execute(select(User).where(User.id == user_id))
    user = result.scalar_one_or_none()

    if user is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="用户不存在"
        )

    if getattr(user, "status", None) == "inactive":
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="账户已被禁用"
        )

    return SuccessResponse(
        data={
            "id": str(user.id),
            "email": user.email,
            "name": user.name or user.full_name,
            "role": user.role,
            "channelId": str(user.channel_id) if user.channel_id is not None else None,
            "status": user.status,
            "subscriptionTier": getattr(user, "subscription_tier", None),
            "lastLoginAt": user.last_login_at.isoformat() if user.last_login_at is not None else None,
        }
    )


@router.post("/logout", response_model=SuccessResponse)
async def logout(
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    用户登出
    
    将当前Token加入黑名单
    """
    from datetime import datetime, timedelta
    from app.token_blacklist import add_token_to_blacklist
    import uuid as uuid_module
    
    user_id = principal.get("user_id")
    claims = principal.get("claims", {})
    
    # 确保 user_id 是有效的 UUID 字符串
    if user_id:
        try:
            # 验证并转换为标准 UUID 字符串格式
            user_id = str(uuid_module.UUID(str(user_id)))
        except (ValueError, TypeError):
            # 如果无法转换，使用 claims 中的 sub
            sub = claims.get("sub")
            if sub:
                try:
                    user_id = str(uuid_module.UUID(str(sub)))
                except (ValueError, TypeError):
                    user_id = None
    
    if not user_id:
        # 如果无法获取有效的 user_id，直接返回成功（前端会清除本地 token）
        return SuccessResponse(message="登出成功")
    
    # 获取Token的过期时间（从claims中提取）
    exp = claims.get("exp")
    if exp:
        expires_at = datetime.utcfromtimestamp(exp)
    else:
        # 默认7天后过期
        expires_at = datetime.utcnow() + timedelta(days=7)
    
    # 生成token的唯一标识（使用sub + 时间戳）
    token_jti = f"logout_{user_id}_{datetime.utcnow().timestamp()}"
    
    try:
        # 添加到黑名单
        await add_token_to_blacklist(
            token_jti=token_jti,
            user_id=user_id,
            expires_at=expires_at,
            reason="logout",
            db=db,
        )
    except Exception as e:
        # 即使黑名单添加失败，也返回成功（前端会清除本地 token）
        import structlog
        logger = structlog.get_logger(__name__)
        logger.warning("failed_to_add_token_to_blacklist", error=str(e), user_id=user_id)
    
    return SuccessResponse(message="登出成功")


@router.post("/refresh", response_model=SuccessResponse)
async def refresh_token_endpoint(
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    刷新访问令牌
    
    使用 Refresh Token 获取新的 Access Token
    """
    user_id = principal.get("user_id")
    claims = principal.get("claims", {})
    
    # 验证是否为 refresh token
    if claims.get("type") != "refresh":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="必须使用 refresh token 调用此接口"
        )
    
    # 查询用户信息
    result = await db.execute(select(User).where(User.id == user_id))
    user = result.scalar_one_or_none()
    
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="用户不存在"
        )
    
    # 生成新的 access token 和 refresh token
    token_data = {
        "sub": str(user.id),
        "email": user.email,
        "role": user.role,
        "channelId": str(user.channel_id) if user.channel_id else None,
    }
    new_access_token = create_access_token(data=token_data)
    new_refresh_token = create_refresh_token(data=token_data)
    
    return SuccessResponse(
        data={
            "token": new_access_token,
            "refreshToken": new_refresh_token,
        }
    )


@router.put("/password", response_model=SuccessResponse)
async def change_password(
    req: PasswordChangeRequest,
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    修改密码
    """
    user_id = principal.get("user_id")
    
    # 查询用户
    result = await db.execute(select(User).where(User.id == user_id))
    user = result.scalar_one_or_none()
    
    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="用户不存在"
        )
    
    # 验证旧密码
    password_hash = user.password_hash or user.hashed_password
    if not verify_password(req.old_password, password_hash):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="旧密码错误"
        )
    
    # 更新密码
    new_hash = get_password_hash(req.new_password)
    user.password_hash = new_hash
    user.hashed_password = new_hash  # 兼容旧字段
    
    await db.commit()
    
    return SuccessResponse(message="密码修改成功")


@router.get("/keys/info", response_model=SuccessResponse)
async def get_api_key_info(
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    获取服务终结点和API密钥信息
    """
    user_id = principal.get("user_id")
    
    # 查询用户的API密钥
    result = await db.execute(
        select(APIKey)
        .where(APIKey.user_id == user_id)
        .where(APIKey.is_active == True)
        .order_by(APIKey.created_at.desc())
    )
    api_key = result.scalar_one_or_none()
    
    if not api_key:
        # 创建新的API密钥
        raw_key = f"sk-{secrets.token_urlsafe(32)}"
        key_hash = get_password_hash(raw_key)
        
        new_key = APIKey(
            user_id=user_id,
            api_key_hash=key_hash,
            api_key_prefix=raw_key[:8],
            name="默认密钥",
            key_hash=key_hash,
            prefix=raw_key[:8],
            is_active=True,
        )
        db.add(new_key)
        await db.commit()
        await db.refresh(new_key)
        
        api_key = new_key
        masked_key = _mask_api_key(raw_key)
    else:
        # 隐藏现有密钥
        masked_key = api_key.api_key_prefix + "..." + "xxxx"
    
    endpoint = f"https://api.taiji-ai.com/v1"  # 可以从配置读取
    
    return SuccessResponse(
        data={
            "endpoint": endpoint,
            "apiKey": masked_key,
            "createdAt": api_key.created_at.isoformat(),
            "lastUsed": api_key.last_used.isoformat() if api_key.last_used else None,
        }
    )


@router.post("/keys/regenerate", response_model=SuccessResponse)
async def regenerate_api_key(
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    重新生成API密钥
    """
    user_id = principal.get("user_id")
    
    # 禁用所有旧密钥
    result = await db.execute(
        select(APIKey).where(APIKey.user_id == user_id)
    )
    old_keys = result.scalars().all()
    
    for key in old_keys:
        key.is_active = False
    
    # 生成新密钥
    raw_key = f"sk-{secrets.token_urlsafe(32)}"
    key_hash = get_password_hash(raw_key)
    
    new_key = APIKey(
        user_id=user_id,
        api_key_hash=key_hash,
        api_key_prefix=raw_key[:8],
        name="重新生成的密钥",
        key_hash=key_hash,
        prefix=raw_key[:8],
        is_active=True,
    )
    db.add(new_key)
    await db.commit()
    
    return SuccessResponse(
        data={
            "apiKey": raw_key,  # 完整显示新密钥
            "message": "旧密钥已失效",
        }
    )


@router.get("/keys", response_model=SuccessResponse)
async def list_api_keys(
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    获取用户的 API 密钥列表
    
    返回所有密钥的信息（不包含完整密钥内容，仅显示前缀）
    """
    user_id = principal.get("user_id")
    
    result = await db.execute(
        select(APIKey)
        .where(APIKey.user_id == user_id)
        .order_by(APIKey.created_at.desc())
    )
    keys = result.scalars().all()
    
    return SuccessResponse(
        data={
            "keys": [
                {
                    "id": str(key.id),
                    "name": key.name or "默认密钥",
                    "prefix": key.api_key_prefix + "...",
                    "isActive": key.is_active,
                    "createdAt": key.created_at.isoformat(),
                    "lastUsed": key.last_used.isoformat() if key.last_used else None,
                    "expiresAt": key.expires_at.isoformat() if key.expires_at else None,
                    "totalRequests": key.total_requests or 0,
                }
                for key in keys
            ],
            "total": len(keys)
        },
        message=f"共 {len(keys)} 个 API 密钥"
    )


@router.post("/keys", response_model=SuccessResponse)
async def create_api_key(
    name: str = Query(default="API Key", description="密钥名称"),
    expires_in_days: Optional[int] = Query(default=None, description="过期天数，不填则永不过期"),
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    创建新的 API 密钥
    
    注意：密钥只在创建时显示一次，请妥善保管
    
    参数：
    - name: 密钥名称，用于标识不同用途的密钥
    - expires_in_days: 过期天数，不填则永不过期
    
    返回：
    - key: 完整的 API 密钥（只在创建时返回一次）
    - id: 密钥 ID，用于删除操作
    """
    user_id = principal.get("user_id")
    
    # 生成新密钥
    raw_key = f"sk-{secrets.token_urlsafe(32)}"
    key_hash = get_password_hash(raw_key)
    
    expires_at = None
    if expires_in_days:
        expires_at = datetime.utcnow() + timedelta(days=expires_in_days)
    
    new_key = APIKey(
        user_id=user_id,
        api_key_hash=key_hash,
        api_key_prefix=raw_key[:8],
        name=name,
        key_hash=key_hash,
        prefix=raw_key[:8],
        is_active=True,
        expires_at=expires_at,
        total_requests=0,
    )
    db.add(new_key)
    await db.commit()
    await db.refresh(new_key)
    
    return SuccessResponse(
        data={
            "id": str(new_key.id),
            "name": name,
            "key": raw_key,  # 只在创建时返回完整密钥
            "prefix": raw_key[:8] + "...",
            "expiresAt": expires_at.isoformat() if expires_at else None,
            "createdAt": new_key.created_at.isoformat(),
        },
        message="API 密钥创建成功，请妥善保管，密钥只显示一次"
    )


@router.delete("/keys/{key_id}", response_model=SuccessResponse)
async def delete_api_key(
    key_id: str,
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db)
):
    """
    删除 API 密钥
    
    参数：
    - key_id: 密钥 ID（从列表接口或创建接口获取）
    
    注意：删除后密钥立即失效，无法恢复
    """
    user_id = principal.get("user_id")
    
    try:
        key_uuid = uuid.UUID(key_id)
    except ValueError:
        raise HTTPException(status_code=400, detail="无效的密钥 ID")
    
    result = await db.execute(
        select(APIKey)
        .where(APIKey.id == key_uuid)
        .where(APIKey.user_id == user_id)
    )
    key = result.scalar_one_or_none()
    
    if not key:
        raise HTTPException(status_code=404, detail="密钥不存在或无权删除")
    
    await db.delete(key)
    await db.commit()
    
    return SuccessResponse(
        data={"id": key_id, "name": key.name},
        message=f"API 密钥 '{key.name}' 已删除"
    )


@router.post("/register/send-code", response_model=SuccessResponse)
async def send_verification_code_endpoint(
    email: str = Query(..., description="邮箱地址"),
    db: AsyncSession = Depends(get_db)
):
    """
    发送邮箱验证码
    
    在用户注册前，先调用此接口发送验证码到邮箱
    
    频率限制：同一邮箱60秒内只能发送一次
    """
    # 1. 检查发送频率限制
    can_send, remaining_seconds = await check_rate_limit(email)
    if not can_send:
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail=f"请等待{remaining_seconds}秒后再重新发送验证码"
        )
    
    # 2. 检查邮箱是否已存在
    result = await db.execute(select(User).where(User.email == email))
    existing_user = result.scalar_one_or_none()
    
    if existing_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="该邮箱已被注册"
        )
    
    # 3. 发送验证码
    code = await send_and_store_verification_code(email)
    if not code:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="验证码发送失败，请稍后重试"
        )
    
    return SuccessResponse(
        message="验证码已发送到您的邮箱，请查收"
    )


@router.post("/register", response_model=SuccessResponse)
async def register(req: UserCreate, db: AsyncSession = Depends(get_db)):
    """
    用户注册接口
    
    允许用户自由注册，创建普通用户账户
    注册成功后自动分配默认资源：
    - 分配到 taiji 渠道 (channelId: b415e70b-8d37-481c-b229-bc3b7871607b)
    - 自定义 agent 配额：2 CPU, 2 GB 内存
    - 平台 agent 各1个
    - 供应商所有模型
    - 余额：20元
    """
    # 1. 先检查邮箱是否已存在（不消耗验证码）
    result = await db.execute(select(User).where(User.email == req.email))
    existing_user = result.scalar_one_or_none()
    
    if existing_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="该邮箱已被注册"
        )
    
    # 2. 检查用户名是否已存在
    #    注意：line 778 处会用 `req.username or req.email.split("@")[0]` 兜底，
    #    所以预检必须按真实写入的 username 来查，否则 fallback 出来的 alice
    #    没经预检，在 line 797 flush 才暴露 IntegrityError。
    effective_username = (req.username or "").strip() or req.email.split("@")[0]
    result = await db.execute(select(User).where(User.username == effective_username))
    existing_username = result.scalar_one_or_none()
    if existing_username:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="该用户名已被使用"
        )
    
    # 3. 验证邮箱验证码（先 peek 不消费 —— 等 DB commit + 外部副作用全部成功
    #    后再调 verify_code 真消费。否则后续任何步骤失败，用户的验证码就被
    #    白白烧掉了，必须重新发码。
    is_valid = await peek_verification_code(req.email, req.verification_code)
    if not is_valid:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="验证码错误或已过期"
        )
    
    # taiji 渠道 ID
    TAIJI_CHANNEL_ID = uuid.UUID("b415e70b-8d37-481c-b229-bc3b7871607b")
    
    # 验证渠道存在
    channel_result = await db.execute(
        select(Channel).where(Channel.id == TAIJI_CHANNEL_ID)
    )
    channel = channel_result.scalar_one_or_none()
    if not channel:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="渠道不存在，请联系管理员"
        )
    
    # 创建新用户
    password_hash = get_password_hash(req.password)
    # 用同一套 effective_username 写入 —— line 745 的预检和这里的写入必须 100%
    # 匹配，否则带空格的 username（如 "alice "）会预检通过但写入不一致。
    username = effective_username
    name = req.full_name or username
    
    new_user = User(
        name=name,
        email=req.email,
        password_hash=password_hash,
        hashed_password=password_hash,  # 兼容字段
        username=username,
        full_name=req.full_name or name,
        role="user",  # 默认角色为普通用户
        channel_id=TAIJI_CHANNEL_ID,  # 分配到 taiji 渠道
        status="active",
        is_active=True,
        is_admin=False,
        credit_limit=0,
        total_eu_consumed=0,
    )
    
    try:
        db.add(new_user)
        await db.flush()  # 获取 user.id，这里会触发唯一约束检查
    except IntegrityError as e:
        await db.rollback()
        error_str = str(e.orig) if e.orig else str(e)
        if "email" in error_str.lower() or "users_email_key" in error_str.lower():
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="该邮箱已被注册"
            )
        elif "username" in error_str.lower() or "users_username_key" in error_str.lower():
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="该用户名已被使用"
            )
        else:
            logger.error(f"注册时数据库约束冲突: {e}")
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="注册信息与现有用户冲突，请更换邮箱或用户名"
            )
    
    user_id = new_user.id

    # 跟踪本次注册在 LiteLLM 远端已经创建的 key —— commit 失败时回扫 delete，
    # 避免远端泄漏 orphan key。注意：generate_key 是 fire-and-forget 网络副作
    # 用，DB rollback 不能撤销它，必须显式 delete。
    created_litellm_keys: list[str] = []
    try:
        # 1. 创建余额记录，初始余额 20 元
        balance = Balance(
            user_id=user_id,
            eu_balance=20.0
        )
        db.add(balance)
        
        # 2. 分配自定义 Agent 配额：2 CPU, 2 GB 内存
        custom_agent_quota = TenantCustomAgentQuota(
            tenant_id=user_id,
            cpu_quota=2.0,
            memory_quota=2.0,
            cpu_used=0.0,
            memory_used=0.0,
            agent_count=0
        )
        db.add(custom_agent_quota)
        
        # 3. 获取所有平台 Agent 模板，为每个模板分配 1 个配额
        try:
            agent_manager_client = get_agent_manager_client()
            platform_templates = await agent_manager_client.list_platform_templates()
            
            for template in platform_templates:
                template_name = template.template
                
                # 获取渠道的配额记录（需要同步更新 pod_used）
                channel_quota_result = await db.execute(
                    select(PlatformAgentQuota)
                    .where(
                        and_(
                            PlatformAgentQuota.target_id == TAIJI_CHANNEL_ID,
                            PlatformAgentQuota.target_type == "channel",
                            PlatformAgentQuota.template_name == template_name
                        )
                    )
                    .with_for_update()  # 行锁
                )
                channel_quota = channel_quota_result.scalar_one_or_none()
                
                # 如果渠道没有该模板的配额，跳过
                if not channel_quota:
                    logger.warning(f"渠道没有平台 Agent '{template_name}' 的配额，跳过分配")
                    continue
                
                # 检查渠道剩余配额是否足够
                remaining = channel_quota.pod_quota - (channel_quota.pod_used or 0)
                if remaining < 1:
                    logger.warning(f"渠道平台 Agent '{template_name}' 配额不足，跳过分配")
                    continue
                
                # 创建平台 Agent 配额记录
                platform_quota = PlatformAgentQuota(
                    target_id=user_id,
                    target_type="tenant",
                    template_name=template_name,
                    pod_quota=1,  # 每个平台 agent 分配 1 个配额
                    pod_used=0,
                    allocated_by=None,  # 系统自动分配
                    allocated_at=datetime.utcnow()
                )
                db.add(platform_quota)
                
                # Bug 修复：更新渠道的 pod_used（分配给租户的配额视为渠道已使用的配额）
                channel_quota.pod_used = (channel_quota.pod_used or 0) + 1
                
                # 同时创建 ResourceAllocation 记录（兼容旧逻辑）
                allocation = ResourceAllocation(
                    target_id=user_id,
                    target_type="tenant",
                    resource_type="agent",
                    resource_id=template_name,
                    quantity=1
                )
                db.add(allocation)
        except AgentManagerError as e:
            # AgentManager 不可达是已知运维态，平台 Agent 分配跳过，注册主流程
            # 继续；其他异常（SQLAlchemy / 编程错误等）让外层 try 兜底回滚。
            logger.warning(f"获取平台 Agent 模板失败，跳过平台 Agent 分配: {e}")
        
        # 4. 分配所有供应商模型
        providers_result = await db.execute(
            select(ModelProvider).where(ModelProvider.is_active == True)
        )
        providers = providers_result.scalars().all()
        
        if channel.litellm_team_id:
            litellm_client = get_litellm_client()

            # 跨 provider 去重：DB 唯一约束是 (tenant_id, model_name)，
            # 如果两个 ModelProvider 行的 supported_models 有重叠（例如都包含
            # taiji/gpt-4o-mini），不去重会触发 uq_tenant_model 导致整个事务
            # 失败 → 注册 500。先以 model_name 为键去重，第一个看到的 provider 胜出。
            seen_models: set[str] = set()
            for provider in providers:
                # 为每个供应商的每个模型创建 TenantModelKey
                for model_name in provider.supported_models:
                    if model_name in seen_models:
                        logger.debug(
                            f"模型 {model_name} 已在其他 provider 处理过，跳过"
                        )
                        continue
                    seen_models.add(model_name)
                    try:
                        # 在 LiteLLM 中创建 Key
                        key = await litellm_client.generate_key(
                            team_id=channel.litellm_team_id,
                            models=[model_name],
                            rpm_limit=provider.rpm or 60,
                            tpm_limit=provider.tpm or 10000,
                            max_budget=500.0,  # 默认预算
                            budget_duration="monthly",
                            key_name=f"tenant-{user_id}-{model_name}",
                            metadata={
                                "tenant_id": str(user_id),
                                "tenant_name": name,
                                "channel_id": str(TAIJI_CHANNEL_ID),
                                "channel_name": channel.name,
                                "model": model_name,
                            }
                        )
                        
                        # 远端 key 创建成功 —— 立刻登记到 created_litellm_keys，
                        # 任何后续失败（DB add/flush/commit）都能 best-effort
                        # 删除掉远端 orphan。
                        created_litellm_keys.append(key.key)

                        # 加密存储 Key
                        encrypted_key = litellm_client.encrypt_key(key.key)

                        # 保存到数据库
                        tenant_key = TenantModelKey(
                            tenant_id=user_id,
                            channel_id=TAIJI_CHANNEL_ID,
                            model_name=model_name,
                            litellm_key_id=key.key,
                            litellm_key_hash=encrypted_key,
                            rpm_limit=provider.rpm or 60,
                            tpm_limit=provider.tpm or 10000,
                            max_budget=500.0,
                            budget_duration="monthly",
                            status="active",
                        )
                        db.add(tenant_key)
                        
                        # 同时记录到 ResourceAllocation
                        model_allocation = ResourceAllocation(
                            target_id=user_id,
                            target_type="tenant",
                            resource_type="model",
                            resource_id=str(provider.id),
                            rpm=provider.rpm or 60,
                            tpm=provider.tpm or 10000,
                        )
                        db.add(model_allocation)
                    except LiteLLMClientError as e:
                        # LiteLLM 远端 API 错误：单模型分配失败可接受，跳过该
                        # 模型继续。注意 generate_key 已在 created_litellm_keys
                        # 追加之前抛出，所以这条路径不会有 orphan。
                        logger.warning(f"为模型 {model_name} 创建 LiteLLM Key 失败: {e}")
                        # 继续处理其他模型
                    # SQLAlchemy / 编程错误等其他异常：故意不 catch —— 让外层 try
                    # 接住执行 rollback + LiteLLM orphan 清理，避免静默返回 200
                    # 但只分到一半 key 的「假成功」。
        else:
            logger.warning(f"渠道 {TAIJI_CHANNEL_ID} 未配置 LiteLLM team，跳过模型分配")
        
        # 提交所有更改
        await db.commit()
        await db.refresh(new_user)

        logger.info(
            f"用户注册成功并分配默认资源",
            user_id=str(user_id),
            email=req.email,
            channel_id=str(TAIJI_CHANNEL_ID),
            litellm_key_count=len(created_litellm_keys),
        )

    except Exception as e:
        # 1. DB 回滚（new_user / balance / quota / TenantModelKey 全部撤销）
        await db.rollback()

        # 2. 关键：DB 回滚撤不掉 LiteLLM 远端已创建的 key —— 必须显式逐个 delete。
        #    best-effort：每个 delete 都 try/except，避免一个失败阻塞剩下的清理。
        if created_litellm_keys:
            logger.warning(
                "注册失败 —— 开始清理 LiteLLM 远端 orphan key",
                email=req.email,
                orphan_count=len(created_litellm_keys),
            )
            try:
                client = get_litellm_client()
                for k in created_litellm_keys:
                    try:
                        await client.delete_key(k)
                    except Exception as del_err:
                        logger.error(
                            "清理 orphan LiteLLM key 失败（需人工跟进）",
                            key_prefix=k[:12] if k else None,
                            email=req.email,
                            error=str(del_err),
                        )
            except Exception as cleanup_err:
                # 连 LiteLLM client 都拿不到：把 orphan key 写到日志里，运维兜底
                logger.error(
                    "无法获取 LiteLLM client 清理 orphan key",
                    email=req.email,
                    orphan_keys_first_chars=[k[:12] for k in created_litellm_keys],
                    error=str(cleanup_err),
                )

        # 3. 注意：绝不能把 str(e) 回给客户端 —— SQLAlchemy IntegrityError 的 str
        #    会包含完整 SQL + 全部 parameters，参数里有 LiteLLM 明文 key
        #    (litellm_key_id 列)。日志里完整记录便于排障，响应只给通用错误码。
        logger.error("用户注册失败", email=req.email, exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail={"code": "REGISTER_FAILED",
                    "message": "注册失败，请稍后重试或联系管理员"},
        )

    # DB commit 已成功 —— 最后才真正消费验证码。这一步即使失败也不再回滚
    # （用户已经注册成功，重复消费没意义；Redis 里的过期码不会被复用，因为
    # 同 email 第二次 register 会在 line 737 的 existing_user 检查处 400）。
    #
    # 用 asyncio.shield 保护 verify_code 不被 client disconnect / task cancel
    # 半道杀掉 —— 否则 CancelledError 是 BaseException 直接穿透 except
    # Exception，Redis 里残留 10min TTL 的码（不影响安全，但占资源）。
    import asyncio  # local import 避免 module-level 顺序问题
    try:
        await asyncio.shield(verify_code(req.email, req.verification_code))
    except asyncio.CancelledError:
        # request 已经被取消，shield 内部的 verify_code 还会跑完。重新抛出，
        # 让 FastAPI 走正常取消流程。
        raise
    except Exception as e:
        logger.warning("注册成功后消费验证码失败（不影响注册结果）",
                       email=req.email, error=str(e))
    
    # 创建JWT token，自动登录
    token_data = {
        "sub": str(new_user.id),
        "email": new_user.email,
        "role": new_user.role,
        "user_id": str(new_user.id),
        "channelId": str(TAIJI_CHANNEL_ID),
    }
    
    # 创建 Access Token（短期有效）和 Refresh Token（长期有效）
    access_token = create_access_token(data=token_data)
    refresh_token = create_refresh_token(data=token_data)
    
    return SuccessResponse(
        data={
            "token": access_token,
            "refreshToken": refresh_token,
            "user": {
                "id": str(new_user.id),
                "name": new_user.name,
                "email": new_user.email,
                "username": new_user.username,
                "role": new_user.role,
                "channelId": str(TAIJI_CHANNEL_ID),
            }
        },
        message="注册成功"
    )


@router.post("/forgot-password/send-code", response_model=SuccessResponse)
async def forgot_password_send_code(
    req: ForgotPasswordRequest,
    db: AsyncSession = Depends(get_db)
):
    """
    忘记密码 - 发送验证码
    
    向用户注册邮箱发送密码重置验证码。
    
    频率限制：同一邮箱60秒内只能发送一次
    
    请求体：
    {
        "email": "user@example.com"
    }
    
    响应：
    - 成功：返回验证码已发送的消息
    - 失败：邮箱不存在、发送频率限制等
    """
    # 1. 检查发送频率限制
    can_send, remaining_seconds = await check_rate_limit(req.email)
    if not can_send:
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail=f"请等待{remaining_seconds}秒后再重新发送验证码"
        )
    
    # 2. 检查邮箱是否存在（必须是已注册用户）
    result = await db.execute(select(User).where(User.email == req.email))
    user = result.scalar_one_or_none()
    
    if not user:
        # 为了安全，不直接告知邮箱不存在，但记录日志
        logger.warning("忘记密码请求：邮箱不存在", email=req.email)
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="该邮箱未注册"
        )
    
    # 3. 发送密码重置验证码（使用专门的密码重置邮件模板）
    code = await send_password_reset_code(req.email)
    if not code:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="验证码发送失败，请稍后重试"
        )
    
    logger.info("忘记密码验证码已发送", email=req.email, user_id=str(user.id))
    
    return SuccessResponse(
        message="验证码已发送到您的邮箱，请查收"
    )


@router.post("/forgot-password/reset", response_model=SuccessResponse)
async def forgot_password_reset(
    req: ResetPasswordRequest,
    db: AsyncSession = Depends(get_db)
):
    """
    忘记密码 - 重置密码
    
    验证邮箱验证码并重置密码。
    
    请求体：
    {
        "email": "user@example.com",
        "verification_code": "123456",
        "new_password": "newPassword123"
    }
    
    响应：
    - 成功：密码重置成功
    - 失败：验证码错误、邮箱不存在等
    """
    # 1. 检查邮箱是否存在
    result = await db.execute(select(User).where(User.email == req.email))
    user = result.scalar_one_or_none()
    
    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="该邮箱未注册"
        )
    
    # 2. 验证验证码
    is_valid = await verify_code(req.email, req.verification_code)
    if not is_valid:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="验证码错误或已过期"
        )
    
    # 3. 更新密码
    new_hash = get_password_hash(req.new_password)
    user.password_hash = new_hash
    user.hashed_password = new_hash  # 兼容旧字段
    
    await db.commit()
    
    logger.info("密码重置成功", email=req.email, user_id=str(user.id))

    return SuccessResponse(
        message="密码重置成功，请使用新密码登录"
    )


# ============================================================
# Heicode §7.8.1 — 内部端点：Heicode 后端 syncLocalUserFromAgnet 调用
# 标记用户的 billing_provider（'newapi' / 'litellm'）
# ============================================================


def _verify_internal_service_token(request: Request) -> None:
    """校验内部服务密钥（不是用户 JWT）。失败 401/403/503。"""
    expected = (settings.heicode_internal_service_token or "").strip()
    if not expected:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail={
                "code": "INTERNAL_SERVICE_AUTH_NOT_CONFIGURED",
                "message": "HEICODE_INTERNAL_SERVICE_TOKEN 未配置",
            },
        )
    auth = (request.headers.get("Authorization") or "").strip()
    if not auth.lower().startswith("bearer "):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Authorization: Bearer <internal service token> required",
        )
    presented = auth[7:].strip()
    # 常量时间比较防 timing attack
    if not secrets.compare_digest(presented, expected):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail={"code": "INVALID_INTERNAL_SERVICE_TOKEN",
                    "message": "internal service token mismatch"},
        )


class _BillingProviderUpdate(BaseModel):
    """内部端点请求体。email 或 user_id 二选一。"""
    email: Optional[str] = None
    user_id: Optional[str] = None
    billing_provider: str  # 'newapi' | 'litellm'


@router.put("/internal/billing-provider", response_model=SuccessResponse)
async def set_billing_provider(
    payload: _BillingProviderUpdate,
    request: Request,
    db: AsyncSession = Depends(get_db),
):
    """Heicode §7.8.1 内部端点：标记用户的 billing_provider。

    **不**用 user JWT 鉴权，用 HEICODE_INTERNAL_SERVICE_TOKEN 共享密钥（避免
    给 Heicode 后端发用户 JWT）。

    用途：Heicode 后端 syncLocalUserFromAgnet 同步出新用户时，调本接口把
    User.billing_provider 设为 'newapi'。
    """
    _verify_internal_service_token(request)

    if payload.billing_provider not in ("newapi", "litellm"):
        raise HTTPException(
            status_code=400,
            detail={"code": "INVALID_BILLING_PROVIDER",
                    "message": "billing_provider 必须是 'newapi' 或 'litellm'"},
        )
    if not payload.email and not payload.user_id:
        raise HTTPException(
            status_code=400,
            detail={"code": "MISSING_USER_REF",
                    "message": "email 或 user_id 二选一必填"},
        )

    user = None
    if payload.user_id:
        try:
            uid = uuid.UUID(payload.user_id)
        except Exception:
            raise HTTPException(status_code=400, detail="user_id 不是有效 UUID")
        user = await db.get(User, uid)
    else:
        result = await db.execute(select(User).where(User.email == payload.email))
        user = result.scalar_one_or_none()

    if user is None:
        raise HTTPException(
            status_code=404,
            detail={"code": "USER_NOT_FOUND",
                    "message": f"找不到用户 (email={payload.email}, user_id={payload.user_id})"},
        )

    old = user.billing_provider
    user.billing_provider = payload.billing_provider
    await db.commit()
    await db.refresh(user)

    logger.info(
        "billing_provider_updated",
        user_id=str(user.id),
        email=user.email,
        old_provider=old,
        new_provider=user.billing_provider,
        actor="heicode_backend_internal",
    )

    return SuccessResponse(data={
        "user_id": str(user.id),
        "email": user.email,
        "billing_provider": user.billing_provider,
        "old_billing_provider": old,
    })