xiangqinxiaochengxu/backend/app/services/upload_service.py

from __future__ import annotations

import imghdr
from io import BytesIO
from pathlib import Path
from tempfile import gettempdir
from time import time
from uuid import uuid4

from fastapi import UploadFile
from PIL import Image, ImageFilter
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.user import User


ALLOWED_IMAGE_TYPES = {"jpeg", "png"}
UPLOAD_ROOT = Path(__file__).resolve().parent.parent / "static" / "uploads"
MAX_PROFILE_IMAGES = 6
MAX_ACTIVITY_IMAGE_SIZE = 5 * 1024 * 1024


def _validate_image(content: bytes) -> str:
    image_type = imghdr.what(None, h=content)
    if image_type not in ALLOWED_IMAGE_TYPES:
        raise ValueError("仅支持 jpg/png 图片")
    return image_type


def _write_temp_file(content: bytes, suffix: str) -> str:
    temp_path = Path(gettempdir()) / f"dating_{uuid4().hex}{suffix}"
    temp_path.write_bytes(content)
    return str(temp_path)


def _save_optimized_variants(content: bytes) -> tuple[str, str, str]:
    _validate_image(content)

    with Image.open(BytesIO(content)) as image:
        image = image.convert("RGB")
        image.thumbnail((1280, 1280))

        main_buffer = BytesIO()
        image.save(main_buffer, format="JPEG", quality=88)
        main_path = _write_temp_file(main_buffer.getvalue(), ".jpg")

        blur_image = image.copy().filter(ImageFilter.GaussianBlur(radius=15))
        blur_buffer = BytesIO()
        blur_image.save(blur_buffer, format="JPEG", quality=80)
        blur_path = _write_temp_file(blur_buffer.getvalue(), "_blur.jpg")

    timestamp = str(int(time()))
    return main_path, blur_path, timestamp


def _persist_avatar_files(user_id: int, timestamp: str, main_path: str, blur_path: str) -> tuple[str, str]:
    user_dir = UPLOAD_ROOT / "avatars" / str(user_id)
    user_dir.mkdir(parents=True, exist_ok=True)

    avatar_name = f"{timestamp}.jpg"
    blur_name = f"{timestamp}_blur.jpg"
    avatar_path = user_dir / avatar_name
    blur_target_path = user_dir / blur_name

    avatar_path.write_bytes(Path(main_path).read_bytes())
    blur_target_path.write_bytes(Path(blur_path).read_bytes())

    return avatar_name, blur_name


def _append_profile_image(existing_images: list[str] | None, avatar_url: str) -> list[str]:
    items = list(existing_images or [])
    items = [item for item in items if item and item != avatar_url]
    items.insert(0, avatar_url)
    return items[:MAX_PROFILE_IMAGES]


async def upload_user_avatar(session: AsyncSession, user: User, file: UploadFile) -> dict:
    content = await file.read()
    _validate_upload_content(content)

    main_path, blur_path, timestamp = _save_optimized_variants(content)

    try:
        avatar_name, blur_name = _persist_avatar_files(user.id, timestamp, main_path, blur_path)
        avatar_url = f"/static/uploads/avatars/{user.id}/{avatar_name}"
        blur_url = f"/static/uploads/avatars/{user.id}/{blur_name}"

        user.avatar_url = avatar_url
        user.avatar_blur_url = blur_url
        user.profile_images = _append_profile_image(user.profile_images, avatar_url)

        session.add(user)
        await session.commit()
        await session.refresh(user)

        return {
            "avatar_url": user.avatar_url,
            "avatar_blur_url": user.avatar_blur_url,
            "profile_images": user.profile_images or [],
            "local_files": {
                "avatar": avatar_name,
                "blur": blur_name,
            },
        }
    finally:
        Path(main_path).unlink(missing_ok=True)
        Path(blur_path).unlink(missing_ok=True)


def _validate_upload_content(content: bytes) -> None:
    if not content:
        raise ValueError("上传文件不能为空")
    if len(content) > MAX_ACTIVITY_IMAGE_SIZE:
        raise ValueError("图片大小不能超过 5MB")


def _store_generic_image(content: bytes, folder: str, filename: str) -> str:
    _validate_image(content)
    target_dir = UPLOAD_ROOT / folder
    target_dir.mkdir(parents=True, exist_ok=True)
    target_path = target_dir / filename
    with Image.open(BytesIO(content)) as image:
        image = image.convert("RGB")
        image.thumbnail((1600, 1600))
        image.save(target_path, format="JPEG", quality=88)
    return str(target_path)


async def upload_activity_image(file: UploadFile, *, folder: str = "activities") -> dict:
    content = await file.read()
    _validate_upload_content(content)

    timestamp = str(int(time()))
    filename = f"{timestamp}_{uuid4().hex}.jpg"
    stored_path = _store_generic_image(content, folder, filename)
    url_path = f"/static/uploads/{folder}/{filename}"
    return {
        "image_url": url_path,
        "local_file": filename,
        "stored_path": stored_path,
    }


async def upload_user_profile_image(session: AsyncSession, user: User, file: UploadFile) -> dict:
    content = await file.read()
    _validate_upload_content(content)

    main_path, blur_path, timestamp = _save_optimized_variants(content)

    try:
        avatar_name, _ = _persist_avatar_files(user.id, f"profile_{timestamp}", main_path, blur_path)
        image_url = f"/static/uploads/avatars/{user.id}/{avatar_name}"
        user.profile_images = _append_profile_image(user.profile_images, image_url)
        if not user.avatar_url:
            user.avatar_url = image_url
        if not user.avatar_blur_url:
            user.avatar_blur_url = image_url

        session.add(user)
        await session.commit()
        await session.refresh(user)

        return {
            "image_url": image_url,
            "profile_images": user.profile_images or [],
        }
    finally:
        Path(main_path).unlink(missing_ok=True)
        Path(blur_path).unlink(missing_ok=True)