huihuiSquare/digital-avatar-app/backend/routers/avatars.py

from fastapi import APIRouter, Depends, Body, Header, UploadFile, File
from sqlalchemy.orm import Session
import os
import uuid
import mimetypes

from database import get_db
from routers.knowledge import UPLOAD_DIR
from models import Avatar, KnowledgeDoc, KnowledgeChunk, QAPair, Authorization, User
from responses import ok, fail

router = APIRouter(tags=["分身"])


def _resolve_user(authorization: str | None, db: Session):
    """从 Authorization: Bearer <app_token> 解析当前登录用户"""
    if not authorization:
        return None
    token = authorization.replace("Bearer ", "", 1).replace("bearer ", "", 1).strip()
    return db.query(User).filter(User.app_token == token).first()


@router.get("/avatar")
def list_avatars(page: int = 1, limit: int = 20, authorization: str = Header(None), db: Session = Depends(get_db)):
    # 仅返回当前登录用户自己的分身；未登录返回空，避免看到种子/他人数据
    user = _resolve_user(authorization, db)
    if not user:
        return ok({"data": [], "total": 0})
    q = db.query(Avatar).filter(Avatar.owner_id == user.huihui_user_id)
    total = q.count()
    items = (
        q.order_by(Avatar.created_at.desc())
        .offset((page - 1) * limit)
        .limit(limit)
        .all()
    )
    return ok({"data": [a.to_dict() for a in items], "total": total})


@router.get("/avatar/{avatar_id}")
def get_avatar(avatar_id: str, db: Session = Depends(get_db)):
    a = db.query(Avatar).filter(Avatar.id == avatar_id).first()
    if not a:
        return fail("分身不存在", 404)
    return ok(a.to_dict())


@router.post("/avatar")
def create_avatar(payload: dict = Body(...), authorization: str = Header(None), db: Session = Depends(get_db)):
    user = _resolve_user(authorization, db)
    a = Avatar(
        owner_id=user.huihui_user_id if user else "",
        name=payload.get("name", "未命名分身"),
        display_name=payload.get("displayName", "") or payload.get("display_name", ""),
        description=payload.get("description", ""),
        photo_url=payload.get("photoUrl", "") or payload.get("photo_url", ""),
        emoji=payload.get("emoji", "🤖"),
        status=payload.get("status", "active"),
        token_balance=payload.get("tokenBalance", 0),
        config=payload.get("config", {}) or {},
    )
    db.add(a)
    db.commit()
    db.refresh(a)
    return ok(a.to_dict())


@router.put("/avatar/{avatar_id}")
def update_avatar(avatar_id: str, payload: dict = Body(...), db: Session = Depends(get_db)):
    a = db.query(Avatar).filter(Avatar.id == avatar_id).first()
    if not a:
        return fail("分身不存在", 404)
    mapping = {
        "displayName": "display_name",
        "photoUrl": "photo_url",
        "tokenBalance": "token_balance",
    }
    for key in ("name", "displayName", "description", "photoUrl", "emoji", "status", "tokenBalance", "config"):
        if key in payload:
            col = mapping.get(key, key)
            setattr(a, col, payload[key])
    db.commit()
    db.refresh(a)
    return ok(a.to_dict())


@router.delete("/avatar/{avatar_id}")
def delete_avatar(avatar_id: str, db: Session = Depends(get_db)):
    a = db.query(Avatar).filter(Avatar.id == avatar_id).first()
    if not a:
        return fail("分身不存在", 404)
    # 级联清理关联数据，避免孤儿记录
    db.query(KnowledgeDoc).filter(KnowledgeDoc.avatar_id == avatar_id).delete()
    db.query(KnowledgeChunk).filter(KnowledgeChunk.avatar_id == avatar_id).delete()
    db.query(QAPair).filter(QAPair.avatar_id == avatar_id).delete()
    db.query(Authorization).filter(Authorization.avatar_id == avatar_id).delete()
    db.delete(a)
    db.commit()
    return ok({"success": True})


@router.post("/avatar/{avatar_id}/photo")
async def upload_avatar_photo(
    avatar_id: str,
    file: UploadFile = File(...),
    db: Session = Depends(get_db),
):
    """上传数字分身头像"""
    a = db.query(Avatar).filter(Avatar.id == avatar_id).first()
    if not a:
        return fail("分身不存在", 404)

    # 验证文件类型
    if not file.content_type or not file.content_type.startswith("image/"):
        return fail("仅支持图片文件", 400)

    file_bytes = await file.read()
    if len(file_bytes) > 5 * 1024 * 1024:
        return fail("头像文件不能超过5MB", 400)

    # 保存到 uploads 目录
    ext = mimetypes.guess_extension(file.content_type) or ".jpg"
    filename = f"avatar-{uuid.uuid4().hex}{ext}"
    avatar_dir = os.path.join(UPLOAD_DIR, avatar_id)
    os.makedirs(avatar_dir, exist_ok=True)
    file_path = os.path.join(avatar_dir, filename)

    with open(file_path, "wb") as f:
        f.write(file_bytes)

    # 更新数据库
    photo_url = f"/api/files/{avatar_id}/{filename}"
    a.photo_url = photo_url
    db.commit()
    db.refresh(a)

    return ok(a.to_dict())