utils/video_processor.py

"""Video I/O via FFmpeg: probe, extract frames, trim, reassemble with audio."""
import os
import json
import shutil
import tempfile
import subprocess
from pathlib import Path
from typing import Optional, Generator, Tuple
from dataclasses import dataclass


@dataclass
class VideoInfo:
    width: int
    height: int
    fps: float
    total_frames: int
    duration: float
    has_audio: bool
    codec: str
    filepath: str


class VideoProcessor:
    def __init__(self, config):
        self.config = config
        self._verify_ffmpeg()

    def _verify_ffmpeg(self):
        try:
            r = subprocess.run(["ffmpeg", "-version"], capture_output=True, text=True, timeout=5)
            if r.returncode != 0:
                raise RuntimeError("FFmpeg not working")
        except FileNotFoundError:
            raise RuntimeError("FFmpeg not found. Install it (packages.txt: ffmpeg).")

    def get_video_info(self, video_path: str) -> VideoInfo:
        cmd = ["ffprobe", "-v", "quiet", "-print_format", "json",
               "-show_format", "-show_streams", video_path]
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            raise ValueError(f"Cannot read video: {video_path}")
        probe = json.loads(r.stdout)
        vstream, has_audio = None, False
        for s in probe.get("streams", []):
            if s["codec_type"] == "video" and vstream is None:
                vstream = s
            elif s["codec_type"] == "audio":
                has_audio = True
        if not vstream:
            raise ValueError("No video stream found")
        fps_parts = vstream.get("r_frame_rate", "30/1").split("/")
        fps = float(fps_parts[0]) / float(fps_parts[1]) if len(fps_parts) == 2 and float(fps_parts[1]) else 30.0
        nb = int(vstream.get("nb_frames", 0))
        dur = float(probe.get("format", {}).get("duration", 0) or 0)
        if nb == 0 and dur > 0:
            nb = int(dur * fps)
        return VideoInfo(int(vstream["width"]), int(vstream["height"]), fps, nb, dur,
                         has_audio, vstream.get("codec_name", "unknown"), video_path)

    def trim(self, video_path: str, start: float, length: float) -> str:
        start = max(0.0, float(start)); length = max(0.1, float(length))
        out_dir = tempfile.mkdtemp(prefix="trim_")
        out_path = os.path.join(out_dir, f"trimmed_{Path(video_path).stem}.mp4")
        cmd = ["ffmpeg", "-y", "-ss", str(start), "-i", video_path, "-t", str(length),
               "-c:v", "libx264", "-preset", "veryfast", "-crf", "18",
               "-c:a", "aac", "-movflags", "+faststart", out_path]
        p = subprocess.run(cmd, capture_output=True, text=True)
        if p.returncode != 0 or not os.path.exists(out_path):
            raise RuntimeError((p.stderr or "trim failed")[-400:])
        return out_path

    def _tmp(self, name):
        d = os.path.join(self.config.video.temp_dir, name)
        os.makedirs(d, exist_ok=True)
        return d

    def extract_frames_to_dir(self, video_path: str, output_dir: str) -> Tuple[str, VideoInfo]:
        os.makedirs(output_dir, exist_ok=True)
        info = self.get_video_info(video_path)
        ext = self.config.video.frame_format
        cmd = ["ffmpeg", "-y", "-i", video_path]
        if ext == "png":
            cmd += ["-compression_level", "3"]
        else:
            cmd += ["-qscale:v", "2"]
        cmd += [os.path.join(output_dir, f"%06d.{ext}")]
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        _, err = p.communicate()
        if p.returncode != 0:
            raise RuntimeError(f"Frame extraction failed: {err.decode()[-400:]}")
        files = sorted(f for f in os.listdir(output_dir) if f.endswith(ext))
        info.total_frames = len(files)
        return output_dir, info

    def assemble_video(self, frames_dir: str, output_path: str, fps: float,
                       original_video: Optional[str] = None) -> str:
        ext = self.config.video.frame_format
        temp_video = output_path + ".temp.mp4"
        cmd = ["ffmpeg", "-y", "-framerate", str(fps),
               "-i", os.path.join(frames_dir, f"%06d.{ext}"),
               "-c:v", self.config.video.output_codec,
               "-crf", str(self.config.video.output_crf),
               "-preset", self.config.video.output_preset,
               "-pix_fmt", self.config.video.pixel_format,
               "-movflags", "+faststart",
               temp_video if original_video else output_path]
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        _, err = p.communicate()
        if p.returncode != 0:
            raise RuntimeError(f"Encoding failed: {err.decode()[-400:]}")
        if original_video:
            info = self.get_video_info(original_video)
            if info.has_audio:
                mux = ["ffmpeg", "-y", "-i", temp_video, "-i", original_video,
                       "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0?",
                       "-shortest", output_path]
                pp = subprocess.Popen(mux, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                _, e2 = pp.communicate()
                if pp.returncode != 0:
                    shutil.move(temp_video, output_path)
                else:
                    os.remove(temp_video)
            else:
                shutil.move(temp_video, output_path)
        return output_path

    def cleanup_temp(self, temp_dir: str):
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)