"""套图任务 API：轮询进度 / 导出 ZIP（进程内内存任务表）。"""
from __future__ import annotations

import io
import zipfile

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse

from schemas import SuiteImageOut, SuiteOut
from services import storage
from services.tasks import IMG_OK, Task, get_task

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


def _task_out(task: Task) -> SuiteOut:
    return SuiteOut(
        id=task.id,
        status=task.status,
        style_set=task.style_set,
        platform=task.platform,
        lang=task.lang,
        ratio=task.ratio,
        provider=task.provider,
        model=task.model,
        total=task.total,
        images=[
            SuiteImageOut(
                type_id=i.type_id, name=i.name, url=i.url,
                status=i.status, error=i.error,
            ) for i in task.images
        ],
        error=task.error,
    )


@router.get("/suites/{suite_id}", response_model=SuiteOut)
async def get_suite(suite_id: str):
    task = get_task(suite_id)
    if task is None:
        raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启），请重新生成")
    return _task_out(task)


@router.get("/suites/{suite_id}/zip")
async def download_suite_zip(suite_id: str):
    """把任务内所有成功图打包成 ZIP（中文文件名）。"""
    task = get_task(suite_id)
    if task is None:
        raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启）")

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        seen: set[str] = set()
        for i, img in enumerate([i for i in task.images if i.status == IMG_OK]):
            path = storage.local_path(img.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\""},
    )