image-suite-studio/server/api/suites.py

"""套图生成 API：创建任务 / 查询状态 / 导出 ZIP。"""
from __future__ import annotations

import io
import zipfile
from uuid import UUID

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from config import get_settings
from db import get_db
from models import Product, ProductAsset, Suite, SuiteImage, STATUS_OK
from schemas import PLATFORM_SPECS, SUPPORTED_TYPES, TONGYI_MODELS, SuiteCreateRequest, SuiteCreateResponse, SuiteImageOut, SuiteOut
from services import storage
from services.generator import run_suite

router = APIRouter(prefix="/api", tags=["suites"])


async def _suite_out(db: AsyncSession, suite: Suite) -> SuiteOut:
    images = (await db.scalars(
        select(SuiteImage).where(SuiteImage.suite_id == suite.id)
        .order_by(SuiteImage.created_at)
    )).all()
    return SuiteOut(
        id=str(suite.id),
        product_id=str(suite.product_id),
        status=suite.status,
        style_set=suite.style_set,
        platform=suite.platform,
        lang=suite.lang,
        ratio=suite.ratio,
        types=list(suite.types or []),
        provider=suite.provider,
        model=suite.model,
        images=[
            SuiteImageOut(
                type_id=i.type_id, name=i.name, url=i.stored_url or "",
                status=i.status, error=i.error,
            ) for i in images
        ],
        error=suite.error,
    )


@router.post("/products/{product_id}/suites", response_model=SuiteCreateResponse)
async def create_suite(
    product_id: str,
    req: SuiteCreateRequest,
    background: BackgroundTasks,
    db: AsyncSession = Depends(get_db),
):
    product = await db.get(Product, UUID(product_id))
    if product is None:
        raise HTTPException(status_code=404, detail="商品不存在")

    # 主图组至少一张图（不要求转存完成：生图可直接用源站 URL 代理解析）
    ok_assets = (await db.scalars(
        select(ProductAsset.id).where(
            ProductAsset.product_id == product.id,
            ProductAsset.group_key == "main",
            ProductAsset.type == "img",
        )
    )).all()
    if not ok_assets:
        raise HTTPException(status_code=400, detail="商品没有主图，无法生成")

    bad = [t for t in req.types if t not in SUPPORTED_TYPES]
    if bad:
        raise HTTPException(status_code=400, detail=f"不支持的图类型: {bad}")
    if req.platform not in PLATFORM_SPECS:
        raise HTTPException(status_code=400, detail=f"不支持的目标平台: {req.platform}（ozon | wb | cn）")
    spec = PLATFORM_SPECS[req.platform]

    settings = get_settings()
    provider_name = req.provider or settings.image_provider
    if provider_name == "tongyi" and req.model and req.model not in TONGYI_MODELS:
        raise HTTPException(status_code=400, detail=f"不支持的模型: {req.model}（tongyi 支持: {TONGYI_MODELS}）")
    suite = Suite(
        product_id=product.id,
        style_set=req.style_set,
        requirements=req.requirements,
        platform=req.platform,
        lang=spec["lang"],
        ratio=spec["ratio"],
        types=req.types,
        provider=provider_name,
        model=req.model,
    )
    db.add(suite)
    await db.commit()
    await db.refresh(suite)

    background.add_task(run_suite, str(suite.id))
    return SuiteCreateResponse(suite_id=str(suite.id))


@router.get("/suites/{suite_id}", response_model=SuiteOut)
async def get_suite(suite_id: str, db: AsyncSession = Depends(get_db)):
    suite = await db.get(Suite, UUID(suite_id))
    if suite is None:
        raise HTTPException(status_code=404, detail="任务不存在")
    return await _suite_out(db, suite)


@router.get("/products/{product_id}/suites")
async def list_suites(product_id: str, db: AsyncSession = Depends(get_db)):
    suites = (await db.scalars(
        select(Suite).where(Suite.product_id == UUID(product_id))
        .order_by(Suite.created_at.desc())
    )).all()
    return [await _suite_out(db, s) for s in suites]


@router.get("/suites/{suite_id}/zip")
async def download_suite_zip(suite_id: str, db: AsyncSession = Depends(get_db)):
    """把任务内所有成功图打包成 ZIP（中文文件名）。"""
    suite = await db.get(Suite, UUID(suite_id))
    if suite is None:
        raise HTTPException(status_code=404, detail="任务不存在")
    images = (await db.scalars(
        select(SuiteImage).where(
            SuiteImage.suite_id == suite.id, SuiteImage.status == STATUS_OK,
        ).order_by(SuiteImage.created_at)
    )).all()

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        seen: set[str] = set()
        for i, img in enumerate(images):
            path = storage.local_path(img.stored_url or "")
            if path is None:
                continue
            filename = img.name or img.type_id
            if filename in seen:  # 同类型多张时加序号防覆盖
                filename = f"{filename}-{i + 1}"
            seen.add(filename)
            zf.write(path, f"{filename}{path.suffix or '.jpg'}")
    buf.seek(0)
    return StreamingResponse(
        buf,
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="suite-{suite_id}.zip"'},
    )