spark-ffmppppeg-api/main.py

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import shutil
import tempfile
from pathlib import Path
from typing import Literal

from fastapi import FastAPI, File, HTTPException, Query, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

logger = logging.getLogger("ffmpeg-api")
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")


FFMPEG_BINARY = os.getenv("FFMPEG_BINARY", "ffmpeg")
ALLOWED_CODECS = {"libx264", "h264_nvenc"}
DEFAULT_CODEC = "h264_nvenc"
DEFAULT_X264_PRESET = "medium"
DEFAULT_NVENC_PRESET = "p4"
FORCE_CUDA = os.getenv("FORCE_CUDA", "").strip().lower() in {"1", "true", "yes", "on"}
OUTPUT_EXTENSION = ".mp4"
CHUNK_SIZE = 1024 * 1024
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_EXTENSIONS = {".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v"}
AUDIO_EXTENSIONS = {".wav", ".mp3", ".m4a", ".aac", ".flac", ".ogg"}


app = FastAPI(
    title="FFmpeg Hardware API",
    description=(
        "API FastAPI per conversione media in MP4/MP3 ottimizzata per CUDA con "
        "`h264_nvenc` come default. Imposta `FORCE_CUDA=1` per accettare solo codec NVIDIA."
    ),
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


CodecType = Literal["libx264", "h264_nvenc"]
ConvertMode = Literal["video_to_mp4", "audio_to_mp3", "audio_image_to_mp4"]


async def save_upload_file(upload_file: UploadFile, destination: Path) -> None:
    with destination.open("wb") as buffer:
        while chunk := await upload_file.read(CHUNK_SIZE):
            buffer.write(chunk)
    await upload_file.close()


def validate_codec(codec: str) -> None:
    if codec not in ALLOWED_CODECS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported codec. Allowed values: {', '.join(sorted(ALLOWED_CODECS))}",
        )
    if FORCE_CUDA and codec != "h264_nvenc":
        raise HTTPException(
            status_code=400,
            detail="FORCE_CUDA is enabled. Only `h264_nvenc` is allowed for video outputs.",
        )


def resolve_preset(codec: str, preset: str | None) -> str:
    if preset:
        return preset
    if codec == "h264_nvenc":
        return DEFAULT_NVENC_PRESET
    return DEFAULT_X264_PRESET


def background_kind(path: Path) -> str:
    suffix = path.suffix.lower()
    if suffix in IMAGE_EXTENSIONS:
        return "image"
    if suffix in VIDEO_EXTENSIONS:
        return "video"
    raise HTTPException(
        status_code=400,
        detail=(
            "Unsupported background format. Use an image "
            f"({', '.join(sorted(IMAGE_EXTENSIONS))}) or a video "
            f"({', '.join(sorted(VIDEO_EXTENSIONS))})."
        ),
    )


def ensure_filename(upload: UploadFile | None, field_name: str) -> str:
    if upload is None or not upload.filename:
        raise HTTPException(status_code=400, detail=f"Missing uploaded file for field '{field_name}'.")
    return upload.filename


def audio_kind(path: Path) -> str:
    suffix = path.suffix.lower()
    if suffix in AUDIO_EXTENSIONS:
        return "audio"
    raise HTTPException(
        status_code=400,
        detail=f"Unsupported audio format. Allowed values: {', '.join(sorted(AUDIO_EXTENSIONS))}.",
    )


async def run_ffmpeg_command(command: list[str], error_prefix: str) -> None:
    process = await asyncio.create_subprocess_exec(
        *command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await process.communicate()

    if process.returncode != 0:
        error_message = (stderr or stdout).decode("utf-8", errors="replace").strip()
        logger.error(
            "%s (ffmpeg rc=%s): %s", error_prefix, process.returncode, error_message or "unknown error"
        )
        raise HTTPException(
            status_code=500,
            detail=f"{error_prefix}: {error_message or 'unknown error'}",
        )


async def run_ffmpeg(
    input_path: Path,
    output_path: Path,
    codec: str,
    preset: str,
) -> None:
    command = [FFMPEG_BINARY, "-y", "-hide_banner", "-loglevel", "error"]

    if codec == "h264_nvenc":
        command.extend(["-hwaccel", "cuda", "-hwaccel_output_format", "cuda"])

    command.extend(
        [
            "-i",
            str(input_path),
            "-map",
            "0:v:0",
            "-map",
            "0:a?",
            "-c:v",
            codec,
            "-preset",
            preset,
            "-pix_fmt",
            "yuv420p",
            "-movflags",
            "+faststart",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            str(output_path),
        ]
    )
    await run_ffmpeg_command(command, "FFmpeg conversion failed")


async def run_ffmpeg_audio_to_mp3(
    input_path: Path,
    output_path: Path,
    bitrate: str,
) -> None:
    audio_kind(input_path)
    command = [
        FFMPEG_BINARY,
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(input_path),
        "-vn",
        "-codec:a",
        "libmp3lame",
        "-b:a",
        bitrate,
        str(output_path),
    ]
    await run_ffmpeg_command(command, "FFmpeg audio conversion failed")


async def run_ffmpeg_compose(
    background_path: Path,
    logo_path: Path,
    audio_path: Path,
    output_path: Path,
    codec: str,
    preset: str,
    width: int,
    height: int,
    fps: int,
    logo_width: int,
    margin: int,
) -> None:
    kind = background_kind(background_path)

    scale_background = (
        f"scale={width}:{height}:force_original_aspect_ratio=decrease,"
        f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2:color=black"
    )
    filter_complex = (
        f"[0:v]{scale_background}[bg];"
        f"[1:v]scale={logo_width}:-1[logo];"
        f"[bg][logo]overlay=W-w-{margin}:{margin}:format=auto[v]"
    )

    command = [
        FFMPEG_BINARY,
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
    ]

    if codec == "h264_nvenc":
        command.extend(["-init_hw_device", "cuda:0"])

    if kind == "image":
        command.extend(["-framerate", str(fps), "-loop", "1"])
    elif codec == "h264_nvenc":
        command.extend(["-hwaccel", "cuda"])

    command.extend(
        [
            "-i",
            str(background_path),
            "-i",
            str(logo_path),
            "-i",
            str(audio_path),
            "-filter_complex",
            filter_complex,
            "-map",
            "[v]",
            "-map",
            "2:a:0",
            "-c:v",
            codec,
            "-preset",
            preset,
            "-r",
            str(fps),
            "-pix_fmt",
            "yuv420p",
            "-movflags",
            "+faststart",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-shortest",
            str(output_path),
        ]
    )

    await run_ffmpeg_command(command, "FFmpeg compose failed")


def cleanup_temp_dir(temp_dir: str) -> None:
    shutil.rmtree(temp_dir, ignore_errors=True)


def output_filename(stem: str, suffix: str) -> str:
    safe_stem = stem or "converted"
    return f"{safe_stem}{suffix}"


@app.get("/health")
async def healthcheck() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/convert")
async def convert_media(
    mode: ConvertMode = Query(
        "video_to_mp4",
        description="Conversion mode: video_to_mp4, audio_to_mp3, audio_image_to_mp4.",
    ),
    file: UploadFile | None = File(None),
    audio: UploadFile | None = File(None),
    background: UploadFile | None = File(None),
    logo: UploadFile | None = File(None),
    codec: CodecType = Query(
        DEFAULT_CODEC,
        description="Codec video consentito. `h264_nvenc` e' il default ottimizzato per CUDA; `libx264` resta disponibile come fallback.",
    ),
    preset: str | None = Query(
        None,
        description="Preset opzionale. Default automatico: `medium` per libx264, `p4` per h264_nvenc.",
    ),
    audio_bitrate: str = Query("192k", description="Audio bitrate for audio_to_mp3 mode."),
    width: int = Query(1920, ge=320, le=7680, description="Final video width for audio_image_to_mp4."),
    height: int = Query(1080, ge=240, le=4320, description="Final video height for audio_image_to_mp4."),
    fps: int = Query(30, ge=1, le=120, description="Frame rate for audio_image_to_mp4."),
    logo_width: int = Query(280, ge=32, le=2048, description="Logo width in pixels."),
    margin: int = Query(24, ge=0, le=500, description="Logo margin from top/right border."),
):
    temp_dir = tempfile.mkdtemp(prefix="ffmpeg-api-")
    uploads = {
        "file": file,
        "audio": audio,
        "background": background,
        "logo": logo,
    }

    try:
        if mode == "video_to_mp4":
            validate_codec(codec)
            resolved_preset = resolve_preset(codec, preset)
            filename = ensure_filename(file, "file")
            input_suffix = Path(filename).suffix or ".bin"
            input_path = Path(temp_dir) / f"input{input_suffix}"
            output_stem = Path(filename).stem or "converted"
            output_path = Path(temp_dir) / f"{output_stem}_{codec}{OUTPUT_EXTENSION}"

            await save_upload_file(file, input_path)
            await run_ffmpeg(input_path=input_path, output_path=output_path, codec=codec, preset=resolved_preset)
            media_type = "video/mp4"

        elif mode == "audio_to_mp3":
            filename = ensure_filename(audio or file, "audio")
            audio_upload = audio or file
            input_suffix = Path(filename).suffix or ".bin"
            input_path = Path(temp_dir) / f"input_audio{input_suffix}"
            output_stem = Path(filename).stem or "converted"
            output_path = Path(temp_dir) / output_filename(output_stem, ".mp3")

            await save_upload_file(audio_upload, input_path)
            await run_ffmpeg_audio_to_mp3(input_path=input_path, output_path=output_path, bitrate=audio_bitrate)
            media_type = "audio/mpeg"

        elif mode == "audio_image_to_mp4":
            validate_codec(codec)
            resolved_preset = resolve_preset(codec, preset)
            audio_filename = ensure_filename(audio, "audio")
            background_filename = ensure_filename(background, "background")
            logo_filename = ensure_filename(logo, "logo")

            background_path = Path(temp_dir) / f"background{Path(background_filename).suffix or '.bin'}"
            logo_path = Path(temp_dir) / f"logo{Path(logo_filename).suffix or '.bin'}"
            audio_path = Path(temp_dir) / f"audio{Path(audio_filename).suffix or '.bin'}"
            output_stem = Path(audio_filename).stem or "composed"
            output_path = Path(temp_dir) / f"{output_stem}_{codec}{OUTPUT_EXTENSION}"

            await save_upload_file(background, background_path)
            await save_upload_file(logo, logo_path)
            await save_upload_file(audio, audio_path)

            await run_ffmpeg_compose(
                background_path=background_path,
                logo_path=logo_path,
                audio_path=audio_path,
                output_path=output_path,
                codec=codec,
                preset=resolved_preset,
                width=width,
                height=height,
                fps=fps,
                logo_width=logo_width,
                margin=margin,
            )
            media_type = "video/mp4"

        else:
            raise HTTPException(status_code=400, detail=f"Unsupported mode: {mode}")

        if not output_path.exists():
            raise HTTPException(status_code=500, detail="FFmpeg did not produce any output file.")

        return FileResponse(
            path=output_path,
            media_type=media_type,
            filename=output_path.name,
            background=BackgroundTask(cleanup_temp_dir, temp_dir),
        )
    except HTTPException:
        cleanup_temp_dir(temp_dir)
        raise
    except Exception as exc:
        cleanup_temp_dir(temp_dir)
        raise HTTPException(status_code=500, detail=f"Unexpected error: {exc}") from exc
    finally:
        for upload in uploads.values():
            if upload is not None:
                with contextlib.suppress(Exception):
                    await upload.close()