Pre_Seg_Server/backend/routers/media.py

"""Media upload and parsing endpoints."""

import logging
import re
from pathlib import Path
from typing import List, Optional

from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile, status
from sqlalchemy.orm import Session

from database import get_db
from minio_client import upload_file, get_presigned_url
from models import ProcessingTask, Project, User
from progress_events import publish_task_progress_event
from routers.auth import require_editor
from schemas import ProcessingTaskOut
from statuses import PROJECT_STATUS_PARSING, PROJECT_STATUS_PENDING, TASK_STATUS_QUEUED
from worker_tasks import parse_project_media

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/media", tags=["Media"])

ALLOWED_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".png", ".jpg", ".jpeg", ".dcm"}


def natural_filename_key(filename: str) -> tuple[object, ...]:
    return tuple(
        int(part) if part.isdigit() else part.casefold()
        for part in re.split(r"(\d+)", Path(filename).name)
    )


def _get_ext(filename: str) -> str:
    return Path(filename).suffix.lower()


@router.post(
    "/upload",
    status_code=status.HTTP_201_CREATED,
    summary="Upload a media file",
)
async def upload_media(
    file: UploadFile = File(...),
    project_id: Optional[int] = Form(None),
    db: Session = Depends(get_db),
    current_user: User = Depends(require_editor),
) -> dict:
    """Accept a video, image, or DICOM file and store it in MinIO.

    If project_id is provided, the video_path of the project is updated.
    Returns the presigned URL of the uploaded object.
    """
    if not file.filename:
        raise HTTPException(status_code=400, detail="Missing filename")

    ext = _get_ext(file.filename)
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type: {ext}",
        )

    data = await file.read()
    object_name = f"uploads/{project_id or 'general'}/{file.filename}"

    try:
        upload_file(object_name, data, content_type=file.content_type or "application/octet-stream", length=len(data))
    except Exception as exc:  # noqa: BLE001
        logger.error("Upload failed: %s", exc)
        raise HTTPException(status_code=500, detail="Upload to storage failed") from exc

    file_url = get_presigned_url(object_name, expires=3600)

    if project_id:
        project = db.query(Project).filter(Project.id == project_id).first()
        if not project:
            raise HTTPException(status_code=404, detail="Project not found")
        project.video_path = object_name
        db.commit()
        logger.info("Linked upload to project_id=%s", project_id)
    else:
        # Auto-create a project named after the file
        project = Project(
            name=file.filename,
            description="Auto-created from upload",
            status=PROJECT_STATUS_PENDING,
            video_path=object_name,
            source_type="video",
            owner_user_id=current_user.id,
        )
        db.add(project)
        db.commit()
        db.refresh(project)
        project_id = project.id
        object_name = f"uploads/{project_id}/{file.filename}"
        # Re-upload with corrected path
        upload_file(object_name, data, content_type=file.content_type or "application/octet-stream", length=len(data))
        project.video_path = object_name
        db.commit()
        logger.info("Auto-created project id=%s for upload %s", project_id, file.filename)

    logger.info("Upload complete: %s (size=%d bytes). Async parsing queued.", object_name, len(data))

    return {
        "object_name": object_name,
        "file_url": file_url,
        "size": len(data),
        "project_id": project_id,
        "message": "Upload successful. Parsing job queued.",
    }


@router.post(
    "/upload/dicom",
    status_code=status.HTTP_201_CREATED,
    summary="Upload multiple DICOM files",
)
async def upload_dicom_batch(
    files: List[UploadFile] = File(...),
    project_id: Optional[int] = Form(None),
    db: Session = Depends(get_db),
    current_user: User = Depends(require_editor),
) -> dict:
    """Upload multiple .dcm files for a DICOM series.

    If project_id is provided, files are added to the existing project.
    Otherwise a new DICOM project is created.
    """
    if not files:
        raise HTTPException(status_code=400, detail="No files uploaded")

    sorted_files = sorted(
        [file for file in files if file.filename and file.filename.lower().endswith(".dcm")],
        key=lambda file: natural_filename_key(file.filename or ""),
    )
    if not sorted_files:
        raise HTTPException(status_code=400, detail="No valid DICOM files uploaded")
    uploaded = []

    if project_id:
        project = db.query(Project).filter(Project.id == project_id).first()
        if not project:
            raise HTTPException(status_code=404, detail="Project not found")
    else:
        # Create new DICOM project
        first_name = sorted_files[0].filename or "DICOM_Series"
        project = Project(
            name=first_name,
            description=f"DICOM series with {len(sorted_files)} files",
            status=PROJECT_STATUS_PENDING,
            source_type="dicom",
            owner_user_id=current_user.id,
        )
        db.add(project)
        db.commit()
        db.refresh(project)
        project_id = project.id
        logger.info("Auto-created DICOM project id=%s", project_id)

    for file in sorted_files:
        data = await file.read()
        object_name = f"uploads/{project_id}/dicom/{file.filename}"
        try:
            upload_file(object_name, data, content_type="application/dicom", length=len(data))
            uploaded.append(object_name)
        except Exception as exc:  # noqa: BLE001
            logger.error("Failed to upload DICOM %s: %s", file.filename, exc)

    project.video_path = f"uploads/{project_id}/dicom"
    db.commit()

    return {
        "project_id": project_id,
        "uploaded_count": len(uploaded),
        "message": f"Uploaded {len(uploaded)} DICOM files. Parsing job queued.",
    }


@router.post(
    "/parse",
    status_code=status.HTTP_202_ACCEPTED,
    response_model=ProcessingTaskOut,
    summary="Trigger frame extraction",
)
def parse_media(
    project_id: int,
    source_type: Optional[str] = None,
    parse_fps: Optional[float] = Query(None, gt=0, le=120),
    max_frames: Optional[int] = Query(None, gt=0),
    target_width: int = Query(640, ge=64, le=4096),
    db: Session = Depends(get_db),
    current_user: User = Depends(require_editor),
) -> ProcessingTask:
    """Create a background task for media frame extraction.

    The Celery worker performs the heavy FFmpeg/OpenCV/pydicom work and
    updates the persisted task record as it progresses.
    """
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    if not project.video_path:
        raise HTTPException(status_code=400, detail="Project has no media uploaded")

    effective_source = source_type or project.source_type or "video"
    effective_parse_fps = parse_fps or project.parse_fps or 30.0
    task = ProcessingTask(
        task_type=f"parse_{effective_source}",
        status=TASK_STATUS_QUEUED,
        progress=0,
        message="解析任务已入队",
        project_id=project_id,
        payload={
            "source_type": effective_source,
            "parse_fps": effective_parse_fps,
            "max_frames": max_frames,
            "target_width": target_width,
        },
    )
    project.parse_fps = effective_parse_fps
    project.status = PROJECT_STATUS_PARSING
    db.add(task)
    db.commit()
    db.refresh(task)
    publish_task_progress_event(task)

    async_result = parse_project_media.delay(task.id)
    task.celery_task_id = async_result.id
    db.commit()
    db.refresh(task)

    logger.info("Queued parse task id=%s project_id=%s celery_id=%s", task.id, project_id, async_result.id)
    return task