api/upload.py

import os
import shutil
import logging
import tempfile
from fastapi import APIRouter, HTTPException, Query, UploadFile, File
from pydantic import BaseModel

from services.auth_service import get_current_user
from services.database_service import create_video, update_video_status, get_video
from services.storage_service import get_local_upload_path

logger = logging.getLogger(__name__)

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

CHUNK_TMP_DIR = os.path.join(tempfile.gettempdir(), "uploads")


class UploadRequest(BaseModel):
    filename: str
    mode: str = "balanced"
    file_size: int = 0
    person_query: str = ""


@router.post("/upload")
async def api_upload(req: UploadRequest, token: str = Query(None)):
    if not token:
        raise HTTPException(status_code=401, detail="缺少认证 token")
    try:
        user = await get_current_user(token)
    except ValueError as e:
        raise HTTPException(status_code=401, detail=str(e))

    ext = os.path.splitext(req.filename)[1].lower()
    if ext not in [".mp4", ".mov", ".avi", ".mkv", ".webm"]:
        raise HTTPException(status_code=400, detail="不支持的视频格式")

    if req.file_size > 1024 * 1024 * 1024:
        raise HTTPException(status_code=400, detail="文件大小超过1GB限制")

    video_id = await create_video(
        user_id=user["user_id"],
        original_filename=req.filename,
        mode=req.mode,
        file_size_bytes=req.file_size,
        person_query=req.person_query,
    )

    return {
        "video_id": video_id,
        "storage_mode": "local",
    }


@router.post("/upload/{video_id}/file")
async def api_upload_file(
    video_id: str,
    token: str = Query(None),
    file: UploadFile = File(...),
):
    if not token:
        raise HTTPException(status_code=401, detail="缺少认证 token")
    try:
        user = await get_current_user(token)
    except ValueError as e:
        raise HTTPException(status_code=401, detail=str(e))

    video = await get_video(video_id)
    if not video:
        raise HTTPException(status_code=404, detail="视频不存在")
    if video["user_id"] != user["user_id"]:
        raise HTTPException(status_code=403, detail="无权操作")

    local_path = get_local_upload_path(user["user_id"], video_id, video["original_filename"] or "video.mp4")

    chunk_size = 1024 * 1024
    with open(local_path, "wb") as f:
        while True:
            chunk = await file.read(chunk_size)
            if not chunk:
                break
            f.write(chunk)

    await update_video_status(video_id, "uploaded")

    return {"video_id": video_id, "status": "uploaded", "message": "上传成功"}


class ChunkCompleteRequest(BaseModel):
    total_chunks: int
    filename: str


def _get_chunk_dir(video_id: str) -> str:
    return os.path.join(CHUNK_TMP_DIR, video_id)


@router.post("/upload/{video_id}/chunk")
async def api_upload_chunk(
    video_id: str,
    token: str = Query(None),
    chunk_index: int = Query(...),
    total_chunks: int = Query(...),
    chunk: UploadFile = File(...),
):
    if not token:
        raise HTTPException(status_code=401, detail="缺少认证 token")
    try:
        user = await get_current_user(token)
    except ValueError as e:
        raise HTTPException(status_code=401, detail=str(e))

    video = await get_video(video_id)
    if not video:
        raise HTTPException(status_code=404, detail="视频不存在")
    if video["user_id"] != user["user_id"]:
        raise HTTPException(status_code=403, detail="无权操作")

    if chunk_index < 0 or chunk_index >= total_chunks:
        raise HTTPException(status_code=400, detail="分片索引越界")

    chunk_dir = _get_chunk_dir(video_id)
    os.makedirs(chunk_dir, exist_ok=True)

    chunk_path = os.path.join(chunk_dir, f"chunk_{chunk_index:04d}")
    chunk_size = 1024 * 1024
    try:
        with open(chunk_path, "wb") as f:
            while True:
                data = await chunk.read(chunk_size)
                if not data:
                    break
                f.write(data)
    except Exception as e:
        logger.exception("保存分片失败 video_id=%s chunk_index=%s", video_id, chunk_index)
        raise HTTPException(status_code=500, detail=f"保存分片失败: {e}")

    # 移除 listdir 调用以减少 IO 开销；前端通过 uploadedChunks 自行跟踪进度
    return {
        "video_id": video_id,
        "chunk_index": chunk_index,
        "received": chunk_index + 1,
        "total_chunks": total_chunks,
    }


@router.post("/upload/{video_id}/complete")
async def api_upload_complete(
    video_id: str,
    token: str = Query(None),
    req: ChunkCompleteRequest = ...,
):
    if not token:
        raise HTTPException(status_code=401, detail="缺少认证 token")
    try:
        user = await get_current_user(token)
    except ValueError as e:
        raise HTTPException(status_code=401, detail=str(e))

    video = await get_video(video_id)
    if not video:
        raise HTTPException(status_code=404, detail="视频不存在")
    if video["user_id"] != user["user_id"]:
        raise HTTPException(status_code=403, detail="无权操作")

    chunk_dir = _get_chunk_dir(video_id)
    if not os.path.isdir(chunk_dir):
        raise HTTPException(status_code=400, detail="未找到分片数据，请重新上传")

    expected = req.total_chunks
    chunk_files = []
    for i in range(expected):
        p = os.path.join(chunk_dir, f"chunk_{i:04d}")
        if not os.path.exists(p):
            raise HTTPException(status_code=400, detail=f"分片 {i} 缺失，请重新上传该分片")
        chunk_files.append(p)

    original_filename = req.filename or video.get("original_filename") or "video.mp4"
    final_path = get_local_upload_path(user["user_id"], video_id, original_filename)

    try:
        with open(final_path, "wb") as out:
            for p in chunk_files:
                with open(p, "rb") as part:
                    shutil.copyfileobj(part, out)
    except Exception as e:
        logger.exception("合并分片失败 video_id=%s", video_id)
        raise HTTPException(status_code=500, detail=f"合并分片失败: {e}")

    final_size = os.path.getsize(final_path)
    expected_size = video.get("file_size_bytes") or 0
    if expected_size and final_size != expected_size:
        logger.warning(
            "合并后大小不一致 video_id=%s expected=%s actual=%s",
            video_id, expected_size, final_size,
        )

    shutil.rmtree(chunk_dir, ignore_errors=True)

    await update_video_status(video_id, "uploaded")

    return {
        "video_id": video_id,
        "status": "uploaded",
        "message": "上传成功",
        "file_size": final_size,
    }


@router.delete("/upload/{video_id}/chunk")
async def api_upload_cancel(
    video_id: str,
    token: str = Query(None),
):
    if not token:
        raise HTTPException(status_code=401, detail="缺少认证 token")
    try:
        user = await get_current_user(token)
    except ValueError as e:
        raise HTTPException(status_code=401, detail=str(e))

    video = await get_video(video_id)
    if not video:
        raise HTTPException(status_code=404, detail="视频不存在")
    if video["user_id"] != user["user_id"]:
        raise HTTPException(status_code=403, detail="无权操作")

    chunk_dir = _get_chunk_dir(video_id)
    if os.path.isdir(chunk_dir):
        shutil.rmtree(chunk_dir, ignore_errors=True)

    return {"video_id": video_id, "status": "cancelled", "message": "已清理分片"}