xiaoxia-saas/apps/api/app/api/routes/chunked_upload.py

"""
Chunked upload routes for large file uploads (up to 2GB).
Supports chunked upload, resume, and automatic cleanup of expired uploads.
"""

import fcntl
import json
import logging
import shutil
import tempfile
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

from app.api.routes._helpers import require_project_and_library
from app.auth import AuthenticatedUser, get_current_user
from app.core.celery_app import celery_app
from app.core.storage import OSSStorageService, get_storage_service
from app.dependencies import (
    get_asset_library_repository,
    get_asset_repository,
    get_ingest_job_repository,
    get_project_repository,
)
from app.schemas.chunked_upload import (
    ChunkedUploadCompleteRequest,
    ChunkedUploadCompleteResponse,
    ChunkedUploadInitRequest,
    ChunkedUploadInitResponse,
    ChunkedUploadStatusResponse,
)
from fastapi import APIRouter, Depends, HTTPException, UploadFile, status
from fastapi.params import File

from packages.application import GetProjectUseCase, SubmitIngestJobCommand, SubmitIngestJobUseCase

router = APIRouter()
logger = logging.getLogger(__name__)

# Configuration
DEFAULT_CHUNK_SIZE = 5 * 1024 * 1024  # 5MB
MAX_FILE_SIZE = 2 * 1024 * 1024 * 1024  # 2GB
CHUNK_EXPIRY_HOURS = 24

# Allowed file types — must stay in sync with upload.py ALLOWED_MIME_TYPES
ALLOWED_MIME_TYPES = {
    # Images
    "image/jpeg",
    "image/png",
    "image/gif",
    "image/webp",
    "image/bmp",
    "image/tiff",
    "image/svg+xml",
    # Video
    "video/mp4",
    "video/quicktime",
    "video/mpeg",
    "video/x-msvideo",
    "video/webm",
    "video/x-matroska",
    "video/3gpp",
    # Audio
    "audio/mpeg",
    "audio/wav",
    "audio/ogg",
    "audio/mp3",
    "audio/flac",
    "audio/aac",
    "audio/x-m4a",
    "audio/webm",
}

# Chunk storage root directory
CHUNK_STORAGE_ROOT = Path(tempfile.gettempdir()) / "chunked_uploads"


def _get_chunk_dir(upload_id: str) -> Path:
    """Get chunk storage directory"""
    return CHUNK_STORAGE_ROOT / upload_id


def _get_upload_meta_path(upload_id: str) -> Path:
    """Get upload metadata file path"""
    return CHUNK_STORAGE_ROOT / f"{upload_id}.meta.json"


def _atomic_check_and_record(upload_id: str, chunk_index: int) -> bool:
    """
    Atomically check if chunk is uploaded and record if not.
    Uses file locking to prevent race conditions.

    Returns:
        True if chunk was newly recorded, False if already exists
    """
    meta_path = _get_upload_meta_path(upload_id)
    CHUNK_STORAGE_ROOT.mkdir(parents=True, exist_ok=True)

    with open(meta_path, "r+", encoding="utf-8") as f:
        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
        try:
            meta = json.load(f)
            if chunk_index in meta["uploaded_chunks"]:
                return False
            meta["uploaded_chunks"].append(chunk_index)
            meta["status"] = "uploading"
            f.seek(0)
            json.dump(meta, f, ensure_ascii=False, indent=2)
            f.truncate()
            return True
        finally:
            fcntl.flock(f.fileno(), fcntl.LOCK_UN)


def _load_upload_meta(upload_id: str) -> dict[str, Any]:
    """Load upload metadata"""
    meta_path = _get_upload_meta_path(upload_id)
    if not meta_path.exists():
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Upload not found")

    with open(meta_path, "r", encoding="utf-8") as f:
        return json.load(f)


def _save_upload_meta(upload_id: str, meta: dict[str, Any]) -> None:
    """Save upload metadata"""
    meta_path = _get_upload_meta_path(upload_id)
    CHUNK_STORAGE_ROOT.mkdir(parents=True, exist_ok=True)
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)


def _validate_file_type(content: bytes, filename: str) -> str:
    """Validate file type"""
    try:
        import magic

        detected_mime = magic.from_buffer(content, mime=True)
    except ImportError:
        import mimetypes

        detected_mime = mimetypes.guess_type(filename)[0] or "application/octet-stream"

    if detected_mime not in ALLOWED_MIME_TYPES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Unsupported file type: {detected_mime}. Allowed types: {', '.join(sorted(ALLOWED_MIME_TYPES))}",
        )
    return detected_mime


def _cleanup_expired_uploads() -> int:
    """Cleanup expired uploads, returns number of cleaned uploads"""
    if not CHUNK_STORAGE_ROOT.exists():
        return 0

    now = datetime.now(timezone.utc)
    cleaned = 0

    for meta_file in CHUNK_STORAGE_ROOT.glob("*.meta.json"):
        try:
            with open(meta_file, "r", encoding="utf-8") as f:
                meta = json.load(f)

            expires_at = datetime.fromisoformat(meta["expires_at"])
            if expires_at.tzinfo is None:
                expires_at = expires_at.replace(tzinfo=timezone.utc)

            # Only cleanup uploads that are not actively being uploaded
            if expires_at < now and meta.get("status") != "uploading":
                upload_id = meta["upload_id"]
                chunk_dir = _get_chunk_dir(upload_id)
                if chunk_dir.exists():
                    shutil.rmtree(chunk_dir)
                meta_file.unlink()
                cleaned += 1
                logger.info(f"Cleaned up expired upload: {upload_id}")
        except Exception as e:
            logger.warning(f"Failed to cleanup upload metadata {meta_file}: {e}")

    return cleaned


@router.post("/init", response_model=ChunkedUploadInitResponse)
async def init_chunked_upload(
    request: ChunkedUploadInitRequest,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    project_repository: Any = Depends(get_project_repository),
    asset_library_repository: Any = Depends(get_asset_library_repository),
) -> ChunkedUploadInitResponse:
    """Initialize chunked upload"""

    # Validate file size
    if request.file_size > MAX_FILE_SIZE:
        raise HTTPException(
            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
            detail=f"File exceeds maximum size ({MAX_FILE_SIZE // (1024 * 1024 * 1024)}GB)",
        )

    # Validate project exists
    project = GetProjectUseCase(project_repository).execute(request.project_id)
    if project is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Project not found")

    # Verify asset library
    require_project_and_library(
        request.project_id,
        request.library_id,
        project_repository,
        asset_library_repository,
    )

    # Calculate chunk size
    chunk_size = DEFAULT_CHUNK_SIZE
    expected_chunks = (request.file_size + chunk_size - 1) // chunk_size
    if expected_chunks != request.total_chunks:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"total_chunks mismatch. Expected {expected_chunks} for file size {request.file_size} with chunk size {chunk_size}",
        )

    # Cleanup expired uploads
    _cleanup_expired_uploads()

    # Generate upload ID
    upload_id = uuid4().hex
    now = datetime.now(timezone.utc)
    expires_at = now + timedelta(hours=CHUNK_EXPIRY_HOURS)

    # Create chunk directory
    chunk_dir = _get_chunk_dir(upload_id)
    chunk_dir.mkdir(parents=True, exist_ok=True)

    # Save metadata
    meta = {
        "upload_id": upload_id,
        "filename": request.filename.replace("/", "_").replace("\\", "_"),
        "file_size": request.file_size,
        "total_chunks": request.total_chunks,
        "uploaded_chunks": [],
        "content_type": request.content_type,
        "project_id": request.project_id,
        "library_id": request.library_id,
        "status": "pending",
        "created_at": now.isoformat(),
        "expires_at": expires_at.isoformat(),
    }
    _save_upload_meta(upload_id, meta)

    return ChunkedUploadInitResponse(
        upload_id=upload_id,
        chunk_size=chunk_size,
        total_chunks=request.total_chunks,
        filename=meta["filename"],
        expires_at=expires_at,
    )


@router.get("/{upload_id}/status", response_model=ChunkedUploadStatusResponse)
async def get_upload_status(
    upload_id: str,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
) -> ChunkedUploadStatusResponse:
    """Get upload status (for resume)"""
    meta = _load_upload_meta(upload_id)

    return ChunkedUploadStatusResponse(
        upload_id=upload_id,
        filename=meta["filename"],
        file_size=meta["file_size"],
        total_chunks=meta["total_chunks"],
        uploaded_chunks=sorted(meta["uploaded_chunks"]),
        status=meta["status"],
        created_at=datetime.fromisoformat(meta["created_at"]),
        expires_at=datetime.fromisoformat(meta["expires_at"]),
    )


@router.post("/{upload_id}/complete", response_model=ChunkedUploadCompleteResponse)
async def complete_chunked_upload(
    upload_id: str,
    request: ChunkedUploadCompleteRequest,
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
    project_repository: Any = Depends(get_project_repository),
    asset_library_repository: Any = Depends(get_asset_library_repository),
    asset_repository: Any = Depends(get_asset_repository),
    ingest_job_repository: Any = Depends(get_ingest_job_repository),
    storage_service: OSSStorageService = Depends(get_storage_service),
) -> ChunkedUploadCompleteResponse:
    """Complete chunked upload, merge chunks"""
    # Load metadata
    meta = _load_upload_meta(upload_id)

    # Verify project ID and library ID
    if request.project_id != meta["project_id"] or request.library_id != meta["library_id"]:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Project or library ID mismatch")

    # Verify all chunks are uploaded
    expected_chunks = set(range(meta["total_chunks"]))
    uploaded_chunks = set(meta["uploaded_chunks"])
    missing_chunks = expected_chunks - uploaded_chunks

    if missing_chunks:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Missing chunks: {sorted(missing_chunks)}. Please upload remaining chunks first.",
        )

    # Validate file type
    chunk_dir = _get_chunk_dir(upload_id)
    sample_chunk_path = chunk_dir / "chunk_000000"
    if sample_chunk_path.exists():
        with open(sample_chunk_path, "rb") as f:
            sample_data = f.read(8192)  # Read first 8KB for type detection
        detected_mime = _validate_file_type(sample_data, meta["filename"])
        if detected_mime not in ALLOWED_MIME_TYPES:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Unsupported file type: {detected_mime}",
            )

    # Merge chunks to temp file
    temp_file_path = CHUNK_STORAGE_ROOT / f"{upload_id}_complete.tmp"
    try:
        with open(temp_file_path, "wb") as out_file:
            for i in range(meta["total_chunks"]):
                chunk_path = chunk_dir / f"chunk_{i:06d}"
                with open(chunk_path, "rb") as in_file:
                    shutil.copyfileobj(in_file, out_file)

        # Verify file size
        actual_size = temp_file_path.stat().st_size
        if actual_size != meta["file_size"]:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"File size mismatch. Expected {meta['file_size']}, got {actual_size}",
            )

        # Upload to OSS
        file_id = uuid4().hex[:8]
        safe_filename = meta["filename"]
        storage_key = f"uploads/{file_id}/{safe_filename}"

        file_url = storage_service.upload_file(
            str(temp_file_path),
            storage_key,
            content_type=meta["content_type"],
        )

        # ── 素材去重检测：同素材库 + 同 file_hash 视为重复 ──
        if request.file_hash:
            existing = asset_repository.find_by_library_and_file_hash(
                library_id=request.library_id,
                file_hash=request.file_hash,
            )
            if existing is not None:
                logger.info(
                    "素材去重命中(chunked): library=%s hash=%s existing_asset=%s",
                    request.library_id,
                    request.file_hash,
                    existing.id,
                )
                meta["status"] = "completed"
                _save_upload_meta(upload_id, meta)
                return ChunkedUploadCompleteResponse(
                    storage_key=storage_key,
                    ingest_job_id="",
                    url=file_url,
                    duplicated=True,
                    asset_id=existing.id,
                )

        # Create ingest job
        use_case = SubmitIngestJobUseCase(ingest_job_repository)
        job = use_case.execute(
            SubmitIngestJobCommand(
                project_id=meta["project_id"],
                library_id=meta["library_id"],
                storage_key=storage_key,
                file_hash=request.file_hash,
            )
        )
        celery_app.send_task("worker.ingest_asset", args=[job.id])

        # Update metadata status
        meta["status"] = "completed"
        _save_upload_meta(upload_id, meta)

        return ChunkedUploadCompleteResponse(
            storage_key=storage_key,
            ingest_job_id=job.id,
            url=file_url,
        )

    finally:
        # Cleanup temp file and chunks
        if temp_file_path.exists():
            temp_file_path.unlink()
        if chunk_dir.exists():
            shutil.rmtree(chunk_dir)
        # Delete metadata file
        meta_path = _get_upload_meta_path(upload_id)
        if meta_path.exists():
            meta_path.unlink()


@router.post("/{upload_id}/{chunk_index}")
async def upload_chunk(
    upload_id: str,
    chunk_index: int,
    chunk: UploadFile = File(..., description="Chunk data"),
    authenticated_user: AuthenticatedUser = Depends(get_current_user),
) -> dict[str, Any]:
    """Upload a single chunk"""
    # Load metadata
    meta = _load_upload_meta(upload_id)

    # Check expiry
    expires_at = datetime.fromisoformat(meta["expires_at"])
    if expires_at.tzinfo is None:
        expires_at = expires_at.replace(tzinfo=timezone.utc)

    if expires_at < datetime.now(timezone.utc):
        raise HTTPException(status_code=status.HTTP_410_GONE, detail="Upload has expired")

    # Validate chunk index
    if chunk_index < 0 or chunk_index >= meta["total_chunks"]:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Invalid chunk index. Must be between 0 and {meta['total_chunks'] - 1}",
        )

    # Atomic check and record to prevent race conditions
    if not _atomic_check_and_record(upload_id, chunk_index):
        return {"message": "Chunk already uploaded", "chunk_index": chunk_index}

    # Read chunk data
    chunk_data = await chunk.read()

    # Validate chunk size (last chunk can be smaller than chunk_size)
    expected_size = DEFAULT_CHUNK_SIZE
    if chunk_index == meta["total_chunks"] - 1:
        expected_size = meta["file_size"] - (chunk_index * DEFAULT_CHUNK_SIZE)

    if len(chunk_data) != expected_size:
        # Rollback the recorded chunk
        meta_path = _get_upload_meta_path(upload_id)
        with open(meta_path, "r+", encoding="utf-8") as f:
            fcntl.flock(f.fileno(), fcntl.LOCK_EX)
            try:
                meta = json.load(f)
                if chunk_index in meta["uploaded_chunks"]:
                    meta["uploaded_chunks"].remove(chunk_index)
                f.seek(0)
                json.dump(meta, f, ensure_ascii=False, indent=2)
                f.truncate()
            finally:
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Chunk size mismatch. Expected {expected_size}, got {len(chunk_data)}",
        )

    # Save chunk
    chunk_path = _get_chunk_dir(upload_id) / f"chunk_{chunk_index:06d}"
    with open(chunk_path, "wb") as f:
        f.write(chunk_data)

    # Reload metadata for response
    meta = _load_upload_meta(upload_id)

    return {
        "message": "Chunk uploaded successfully",
        "chunk_index": chunk_index,
        "uploaded_chunks": len(meta["uploaded_chunks"]),
        "total_chunks": meta["total_chunks"],
    }