character-roleplay-backend/app/auth/oauth/google.py

import httpx
from fastapi import APIRouter, Depends, Request, HTTPException, BackgroundTasks
from fastapi.responses import RedirectResponse
import os
import secrets
from app.auth.deps import get_current_user
from app.db import SessionLocal
from app.models import User
from app.auth.jwt import create_jwt_token, create_refresh_token, decode_token
from app.utils.format import get_image_url

router = APIRouter()

GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
GOOGLE_REDIRECT_URI = os.getenv("GOOGLE_REDIRECT_URI")

GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo"

# 定义独立的后台任务函数，用于处理头像转存到 S3 的耗时操作
async def async_upload_avatar_task(google_avatar_url: str, user_id: str):
    """
    后台任务函数：下载 Google 头像并上传到 S3，然后更新用户记录。
    因为主接口的 db 已经关闭，这里需要独立管理 db 的生命周期。
    """
    from app.utils.s3 import upload_google_avatar_to_s3
    
    # 如果 Google 本身就没给头像，无需转存
    if not google_avatar_url:
        return
        
    try:
        # 执行耗时的网络 I/O：下载并上传到 S3
        s3_avatar_url = await upload_google_avatar_to_s3(google_avatar_url, user_id)
        
        # 成功后，开启一个独立的数据库连接写回数据
        db_task = SessionLocal()
        try:
            user = db_task.query(User).filter(User.id == user_id).first()
            if user:
                user.avatar = s3_avatar_url
                db_task.commit()
                print(f"成功在后台为用户 {user_id} 转存头像到 S3")
        finally:
            db_task.close() # 必须手动关闭连接
            
    except Exception as s3_err:
        # 如果 S3 依然失败，打印错误日志，但此时绝不会打扰到已经登录的用户
        print(f"Background S3 Upload failed for user {user_id}: {s3_err}")


@router.get("/login/google")
def login_google():
    state = secrets.token_urlsafe(16)

    url = (
        f"{GOOGLE_AUTH_URL}"
        f"?client_id={GOOGLE_CLIENT_ID}"
        f"&redirect_uri={GOOGLE_REDIRECT_URI}"
        f"&response_type=code"
        f"&scope=openid email profile"
        f"&state={state}"
    )

    response = RedirectResponse(url)
    response.set_cookie(
        "oauth_state",
        state,
        httponly=True,
        secure=True,
        samesite="lax"
    )
    return response

@router.get("/auth/callback")
async def auth_callback(request: Request, code: str, state: str, background_tasks: BackgroundTasks):
    cookie_state = request.cookies.get("oauth_state")

    if not cookie_state or cookie_state != state:
        raise HTTPException(status_code=400, detail="Invalid state")

    async with httpx.AsyncClient() as client:
        token_res = await client.post(
            GOOGLE_TOKEN_URL,
            data={
                "code": code,
                "client_id": GOOGLE_CLIENT_ID,
                "client_secret": GOOGLE_CLIENT_SECRET,
                "redirect_uri": GOOGLE_REDIRECT_URI,
                "grant_type": "authorization_code",
            },
        )

        token_json = token_res.json()
        access_token = token_json["access_token"]

        user_res = await client.get(
            GOOGLE_USERINFO_URL,
            headers={"Authorization": f"Bearer {access_token}"},
        )

        user_info = user_res.json()

    db = SessionLocal()
    
    try:
        user = db.query(User).filter(User.email == user_info["email"]).first()

        if not user:
            # 提取邮箱前缀作为 handle 基础
            # 比如 "john.doe@gmail.com" -> "john_doe"
            base_handle = user_info["email"].split("@")[0].replace(".", "_").lower()
            
            # 拼接一个随机短字符串，确保 handle 的唯一性
            import secrets
            random_suffix = secrets.token_hex(3) # 生成 6 位随机字符
            suggested_handle = f"{base_handle}_{random_suffix}"
                        
            # 如果 Google 没有头像，使用 DiceBear 头像生成器作为兜底
            google_avatar = user_info.get("picture")
            if not google_avatar:
                default_avatar_url = f"https://api.dicebear.com/7.x/lorelei/svg?seed={base_handle}"
            else:
                default_avatar_url = google_avatar
            
            user = User(
                google_id=user_info["id"],
                email=user_info["email"],
                name=user_info["name"],
                avatar=default_avatar_url,
                handle=suggested_handle,
            )
            db.add(user)
            db.commit()
            db.refresh(user)
            
            # 异步/同步将头像转存到 S3
            # FastAPI 会在把下方的 return 响应发送给浏览器之后，立刻在后台启动这个任务
            if google_avatar:
                background_tasks.add_task(
                    async_upload_avatar_task, 
                    google_avatar, 
                    str(user.id)
                )
            
            # 这里不需要再次 refresh，除非后面还要用到更新后的 user 对象
            # db.refresh(user)
    finally:
        db.close()

    jwt_token = create_jwt_token({
        "sub": user.id,
        "email": user.email,
    })
    
    refresh_token = create_refresh_token({
        "sub": user.id,
        "email": user.email,
    })

    return {
        "access_token": jwt_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        # "user": {
        #     "email": user.email,
        #     "name": user.name,
        #     "avatar": get_image_url(user.avatar),
        # }
    }

# @router.get("/me")
# def me(user = Depends(get_current_user)):
#     return {
#         "email": user.email,
#         "name": user.name,
#         "avatar": get_image_url(user.avatar),
#     }
    
# @router.post("/refresh")
# def refresh_token(refresh_token: str):
#     payload = decode_token(refresh_token)

#     if payload.get("type") != "refresh":
#         raise HTTPException(status_code=401)

#     user_id = payload.get("sub")

#     new_access = create_jwt_token({"sub": user_id})

#     return {"access_token": new_access}

# @router.post("/logout")
# def logout():
#     return {"message": "ok"}