src/amuseme/video_assembler.py

"""
video_assembler.py — High-speed video assembly via FFmpeg stdin pipe.
Avoids MoviePy's per-frame overhead by piping raw RGB bytes directly to FFmpeg.
"""
import subprocess
import tempfile
import os
from typing import Generator

from .renderer import WIDTH, HEIGHT, TARGET_FPS
from .logger import get_logger

logger = get_logger("video_assembler")


def assemble(
    frame_generator: Generator[bytes, None, None],
    audio_path: str,
    output_path: str | None = None,
) -> str:
    """
    Pipe frames from frame_generator into FFmpeg, mux with audio_path,
    and write an MP4 to output_path.

    Returns the path to the finished MP4.
    """
    if output_path is None:
        tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        output_path = tmp.name
        tmp.close()

    ffmpeg_cmd = [
        "ffmpeg", "-y",
        # Video input from stdin
        "-f", "rawvideo",
        "-pix_fmt", "rgb24",
        "-s", f"{WIDTH}x{HEIGHT}",
        "-r", str(TARGET_FPS),
        "-i", "pipe:0",
        # Audio input from file
        "-i", audio_path,
        # Output encoding
        "-vcodec", "libx264",
        "-pix_fmt", "yuv420p",
        "-preset", "ultrafast",
        "-crf", "23",
        # Map both streams
        "-map", "0:v:0",
        "-map", "1:a:0",
        "-shortest",
        output_path,
    ]

    logger.info(f"=== ASSEMBLER INPUT === {WIDTH}x{HEIGHT}@{TARGET_FPS}fps, audio={audio_path} → {output_path}")
    process = subprocess.Popen(
        ffmpeg_cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
    )

    try:
        frame_count = 0
        for frame_bytes in frame_generator:
            process.stdin.write(frame_bytes)
            frame_count += 1
            if frame_count % (TARGET_FPS * 5) == 0:
                logger.info(f"[assembler] piped {frame_count} frames ({frame_count / TARGET_FPS:.1f}s)...")

        process.stdin.close()
        process.wait()

        if process.returncode != 0:
            stderr = process.stderr.read().decode(errors="replace")
            raise RuntimeError(f"FFmpeg failed (code {process.returncode}):\n{stderr}")

    except Exception:
        process.kill()
        if os.path.exists(output_path):
            os.remove(output_path)
        raise

    logger.info(f"=== ASSEMBLER OUTPUT === wrote {frame_count} frames → {output_path}")
    return output_path