Pre_Seg_Server/backend/routers/media.py

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

import logging
import os
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Optional

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

from database import get_db
from minio_client import upload_file, get_presigned_url
from models import Project, Frame
from schemas import FrameOut
from services.frame_parser import parse_video, parse_dicom, upload_frames_to_minio

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

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


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),
) -> 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 project:
            project.video_path = object_name
            db.commit()
            logger.info("Linked upload to project_id=%s", project_id)
        else:
            logger.warning("Project id=%s not found for upload linkage", project_id)
    else:
        # Auto-create a project named after the file
        project = Project(
            name=file.filename,
            description="Auto-created from upload",
            status="pending",
            video_path=object_name,
        )
        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)

    # TODO: enqueue async parsing job (Celery / background task)
    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(
    "/parse",
    status_code=status.HTTP_202_ACCEPTED,
    summary="Trigger frame extraction",
)
def parse_media(
    project_id: int,
    source_type: str = "video",  # video | dicom
    db: Session = Depends(get_db),
) -> dict:
    """Trigger frame extraction for a project's uploaded media.

    * video: uses FFmpeg or OpenCV fallback.
    * dicom: uses pydicom to read DCM frames.

    Extracted frames are uploaded to MinIO and registered in the database.
    """
    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")

    # Download from MinIO to a temp directory
    from minio_client import download_file

    try:
        media_bytes = download_file(project.video_path)
    except Exception as exc:  # noqa: BLE001
        logger.error("Failed to download media for parsing: %s", exc)
        raise HTTPException(status_code=500, detail="Failed to retrieve media from storage") from exc

    tmp_dir = tempfile.mkdtemp(prefix=f"seg_parse_{project_id}_")
    local_path = os.path.join(tmp_dir, Path(project.video_path).name)

    with open(local_path, "wb") as f:
        f.write(media_bytes)

    output_dir = os.path.join(tmp_dir, "frames")
    os.makedirs(output_dir, exist_ok=True)

    try:
        if source_type == "dicom":
            # For DICOM, treat local_path as a directory if it contains multiple .dcm
            # If a single .dcm file was uploaded, put it in its own sub-dir
            dcm_dir = os.path.join(tmp_dir, "dcm")
            os.makedirs(dcm_dir, exist_ok=True)
            if local_path.lower().endswith(".dcm"):
                shutil.move(local_path, os.path.join(dcm_dir, os.path.basename(local_path)))
            else:
                shutil.unpack_archive(local_path, dcm_dir) if shutil.which("unzip") else shutil.move(local_path, dcm_dir)
            frame_files = parse_dicom(dcm_dir, output_dir)
        else:
            frame_files = parse_video(local_path, output_dir, fps=30)
    except Exception as exc:  # noqa: BLE001
        logger.error("Frame extraction failed: %s", exc)
        shutil.rmtree(tmp_dir, ignore_errors=True)
        raise HTTPException(status_code=500, detail="Frame extraction failed") from exc

    # Upload frames to MinIO
    try:
        object_names = upload_frames_to_minio(frame_files, project_id)
    except Exception as exc:  # noqa: BLE001
        logger.error("Frame upload failed: %s", exc)
        shutil.rmtree(tmp_dir, ignore_errors=True)
        raise HTTPException(status_code=500, detail="Frame upload to storage failed") from exc

    # Register frames in DB
    frames_out = []
    for idx, obj_name in enumerate(object_names):
        # Get image dimensions
        local_frame = frame_files[idx]
        try:
            import cv2
            img = cv2.imread(local_frame)
            h, w = img.shape[:2] if img is not None else (None, None)
        except Exception:  # noqa: BLE001
            h, w = None, None

        frame = Frame(
            project_id=project_id,
            frame_index=idx,
            image_url=obj_name,
            width=w,
            height=h,
        )
        db.add(frame)
        frames_out.append(frame)

    db.commit()
    for f in frames_out:
        db.refresh(f)

    # Cleanup temp files
    shutil.rmtree(tmp_dir, ignore_errors=True)

    project.status = "ready"
    db.commit()

    logger.info("Parsed %d frames for project_id=%s", len(frames_out), project_id)
    return {
        "project_id": project_id,
        "frames_extracted": len(frames_out),
        "status": "ready",
        "message": "Frame extraction completed successfully.",
    }