video_utils.py

"""ffmpeg/ffprobe wrapper utilities for video processing."""

import json
import subprocess
import tempfile
from pathlib import Path


def get_video_info(video_path: str) -> dict:
    """Get video metadata using ffprobe."""
    path = Path(video_path).expanduser().resolve()
    if not path.exists():
        raise FileNotFoundError(f"Video file not found: {path}")

    cmd = [
        "ffprobe",
        "-v", "quiet",
        "-print_format", "json",
        "-show_format",
        "-show_streams",
        str(path),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    if result.returncode != 0:
        raise RuntimeError(f"ffprobe failed: {result.stderr}")

    probe = json.loads(result.stdout)
    fmt = probe.get("format", {})
    streams = probe.get("streams", [])

    video_stream = next((s for s in streams if s.get("codec_type") == "video"), None)
    audio_stream = next((s for s in streams if s.get("codec_type") == "audio"), None)

    info: dict = {
        "file_path": str(path),
        "file_size_mb": round(int(fmt.get("size", 0)) / (1024 * 1024), 2),
        "duration_seconds": round(float(fmt.get("duration", 0)), 2),
        "format_name": fmt.get("format_name", "unknown"),
        "bit_rate_kbps": round(int(fmt.get("bit_rate", 0)) / 1000),
    }

    if video_stream:
        info["video"] = {
            "codec": video_stream.get("codec_name", "unknown"),
            "width": video_stream.get("width"),
            "height": video_stream.get("height"),
            "fps": _parse_fps(video_stream.get("r_frame_rate", "0/1")),
            "total_frames": int(video_stream.get("nb_frames", 0) or 0),
        }

    if audio_stream:
        info["audio"] = {
            "codec": audio_stream.get("codec_name", "unknown"),
            "sample_rate": audio_stream.get("sample_rate", "unknown"),
            "channels": audio_stream.get("channels", 0),
        }

    return info


def extract_frames(
    video_path: str,
    fps: float = 1.0,
    max_frames: int = 20,
    quality: int = 5,
    output_dir: str | None = None,
) -> dict:
    """Extract frames from video at specified fps.

    Args:
        video_path: Path to video file.
        fps: Frames per second to extract.
        max_frames: Maximum number of frames to extract.
        quality: JPEG quality (2=best, 31=worst).
        output_dir: Directory to save frames. Auto-created if None.

    Returns:
        Dict with output_dir and list of frame paths.
    """
    path = Path(video_path).expanduser().resolve()
    if not path.exists():
        raise FileNotFoundError(f"Video file not found: {path}")

    if output_dir:
        out = Path(output_dir).expanduser().resolve()
        out.mkdir(parents=True, exist_ok=True)
    else:
        out = Path(tempfile.mkdtemp(prefix="video-analyzer-"))

    # Get duration to calculate frame count
    info = get_video_info(str(path))
    duration = info.get("duration_seconds", 0)
    estimated_frames = int(duration * fps)
    actual_max = min(estimated_frames, max_frames) if estimated_frames > 0 else max_frames

    # Use vframes to limit output
    cmd = [
        "ffmpeg",
        "-i", str(path),
        "-vf", f"fps={fps}",
        "-vframes", str(actual_max),
        "-q:v", str(quality),
        "-y",
        str(out / "frame_%04d.jpg"),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg frame extraction failed: {result.stderr}")

    frames = sorted(out.glob("frame_*.jpg"))
    return {
        "output_dir": str(out),
        "frame_count": len(frames),
        "frame_paths": [str(f) for f in frames],
    }


def extract_audio(video_path: str, output_path: str | None = None) -> str:
    """Extract audio track from video as WAV.

    Returns path to extracted audio file.
    """
    path = Path(video_path).expanduser().resolve()
    if not path.exists():
        raise FileNotFoundError(f"Video file not found: {path}")

    if output_path:
        audio_path = Path(output_path).expanduser().resolve()
    else:
        tmp = tempfile.mktemp(prefix="video-analyzer-audio-", suffix=".wav")
        audio_path = Path(tmp)

    cmd = [
        "ffmpeg",
        "-i", str(path),
        "-vn",
        "-acodec", "pcm_s16le",
        "-ar", "16000",
        "-ac", "1",
        "-y",
        str(audio_path),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg audio extraction failed: {result.stderr}")

    return str(audio_path)


def _parse_fps(rate_str: str) -> float:
    """Parse ffprobe frame rate string like '30/1' or '29.97'."""
    try:
        if "/" in rate_str:
            num, den = rate_str.split("/")
            return round(float(num) / float(den), 2) if float(den) != 0 else 0.0
        return round(float(rate_str), 2)
    except (ValueError, ZeroDivisionError):
        return 0.0