backend/scene_processing/frames.py

"""
Frame sampling helpers for the scene pipeline.

For each scene window we want a small set of *representative* frames — enough
for the VLM to read the action, few enough to keep token cost low. The
strategy:

1. Sample N evenly-spaced timestamps inside the window.
2. Extract them with ffmpeg as JPEGs at a reasonable resolution.
3. Drop near-duplicates using a perceptual hash so a static talking head
   doesn't waste 6 frame slots on essentially the same picture.

Each frame is also written to disk under the run's output directory so the
benchmark report can show what the VLM actually saw.
"""

from __future__ import annotations

import os
import shutil
import subprocess
from dataclasses import dataclass
from typing import Sequence

from PIL import Image
import imagehash

from utils import log_message

DEFAULT_FRAMES_PER_WINDOW = 6
DEFAULT_FRAME_HEIGHT = 448
PHASH_DUPLICATE_DISTANCE = 4


@dataclass
class SampledFrame:
    """A frame extracted for a scene window."""

    timestamp_seconds: float
    image_path: str

    def to_dict(self) -> dict:
        return {
            "timestamp_seconds": round(self.timestamp_seconds, 3),
            "image_path": self.image_path,
        }


class FrameExtractionError(RuntimeError):
    """Raised when ffmpeg cannot extract a frame from a video."""


def probe_duration_seconds(video_path: str) -> float:
    """Return the duration of a video in seconds using ffprobe.

    Raises ``FrameExtractionError`` if ffprobe is missing or the video
    cannot be read — never silently returns 0, because a 0 duration would
    cause downstream window construction to produce zero work and the
    operator would have no idea why.
    """
    if shutil.which("ffprobe") is None:
        raise FrameExtractionError(
            "ffprobe is not installed. Install ffmpeg (which includes ffprobe) "
            "before running the scene pipeline."
        )
    if not os.path.exists(video_path):
        raise FrameExtractionError(f"Video file not found: {video_path}")

    result = subprocess.run(
        [
            "ffprobe",
            "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            video_path,
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        raise FrameExtractionError(
            f"ffprobe failed for {video_path}: {result.stderr.strip()}"
        )
    raw = result.stdout.strip()
    if not raw:
        raise FrameExtractionError(f"ffprobe returned no duration for {video_path}")
    try:
        duration = float(raw)
    except ValueError as exc:
        raise FrameExtractionError(
            f"ffprobe returned non-numeric duration {raw!r} for {video_path}"
        ) from exc
    if duration <= 0:
        raise FrameExtractionError(
            f"ffprobe reported non-positive duration {duration} for {video_path}"
        )
    return duration


def sample_window_frames(
    video_path: str,
    window_start: float,
    window_end: float,
    output_dir: str,
    *,
    frames_per_window: int = DEFAULT_FRAMES_PER_WINDOW,
    frame_height: int = DEFAULT_FRAME_HEIGHT,
    dedup_phash: bool = True,
) -> list[SampledFrame]:
    """Extract `frames_per_window` frames from a window and dedup near-duplicates.

    Args:
        video_path: Source mp4 path.
        window_start: Window start in seconds.
        window_end: Window end in seconds (exclusive).
        output_dir: Directory to write extracted JPEGs into. Created if missing.
        frames_per_window: Target frame count before deduplication.
        frame_height: Resize height in pixels. Width is auto (preserves aspect).
        dedup_phash: Drop frames whose perceptual hash is within
            PHASH_DUPLICATE_DISTANCE bits of an earlier kept frame.

    Returns:
        Frames kept after deduplication, ordered by timestamp.
    """
    if shutil.which("ffmpeg") is None:
        raise FrameExtractionError(
            "ffmpeg is not installed. Install ffmpeg before running the scene pipeline."
        )
    if window_end <= window_start:
        return []
    if frames_per_window <= 0:
        return []

    os.makedirs(output_dir, exist_ok=True)
    timestamps = _evenly_spaced_timestamps(window_start, window_end, frames_per_window)

    extracted: list[SampledFrame] = []
    for slot_index, timestamp in enumerate(timestamps):
        image_path = os.path.join(output_dir, f"frame-{slot_index:02d}.jpg")
        try:
            _extract_single_frame(
                video_path=video_path,
                timestamp_seconds=timestamp,
                output_path=image_path,
                frame_height=frame_height,
            )
        except FrameExtractionError as exc:
            log_message(
                f"Skipping frame at {timestamp:.2f}s of {video_path}: {exc}",
                "WARN",
            )
            continue
        extracted.append(SampledFrame(timestamp_seconds=timestamp, image_path=image_path))

    if not dedup_phash or len(extracted) <= 1:
        return extracted

    return _dedup_by_phash(extracted)


def _evenly_spaced_timestamps(start: float, end: float, count: int) -> list[float]:
    """Return N timestamps evenly spaced *inside* (start, end).

    Avoids the exact endpoints — boundary frames are often cuts or fades
    that aren't representative of the window's content.
    """
    if count == 1:
        return [start + (end - start) / 2.0]
    step = (end - start) / (count + 1)
    return [start + step * (i + 1) for i in range(count)]


def _extract_single_frame(
    *,
    video_path: str,
    timestamp_seconds: float,
    output_path: str,
    frame_height: int,
) -> None:
    """Run ffmpeg to grab one frame. Fast-seek before -i for speed."""
    cmd = [
        "ffmpeg",
        "-y",
        "-ss", f"{timestamp_seconds:.3f}",
        "-i", video_path,
        "-frames:v", "1",
        "-vf", f"scale=-2:{frame_height}",
        "-q:v", "3",
        "-loglevel", "error",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    if result.returncode != 0 or not os.path.exists(output_path):
        raise FrameExtractionError(
            f"ffmpeg failed at {timestamp_seconds}s: {result.stderr.strip()}"
        )


def _dedup_by_phash(frames: Sequence[SampledFrame]) -> list[SampledFrame]:
    """Drop near-duplicate frames using perceptual hash distance."""
    kept: list[SampledFrame] = []
    kept_hashes: list[imagehash.ImageHash] = []
    for frame in frames:
        try:
            with Image.open(frame.image_path) as img:
                phash = imagehash.phash(img)
        except (OSError, ValueError) as exc:
            log_message(
                f"phash failed for {frame.image_path}: {exc}; keeping frame anyway",
                "WARN",
            )
            kept.append(frame)
            continue

        is_duplicate = any(
            (phash - prior) <= PHASH_DUPLICATE_DISTANCE for prior in kept_hashes
        )
        if is_duplicate:
            try:
                os.remove(frame.image_path)
            except OSError:
                pass
            continue
        kept.append(frame)
        kept_hashes.append(phash)
    return kept


def extract_audio_clip(
    video_path: str,
    window_start: float,
    window_end: float,
    output_path: str,
    *,
    sample_rate: int = 48000,
) -> str:
    """Extract a window's audio as a mono WAV file for CLAP.

    CLAP expects 48kHz mono. We resample at extraction time so the audio
    tagger doesn't need to do it per-call.

    Returns the output path on success. Raises ``FrameExtractionError``
    on failure rather than returning silently — a missing audio clip
    should be visible in the report, not papered over.
    """
    if shutil.which("ffmpeg") is None:
        raise FrameExtractionError("ffmpeg is not installed.")
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    duration = max(0.0, window_end - window_start)
    if duration <= 0:
        raise FrameExtractionError(
            f"Non-positive window duration ({duration}s) for audio extraction"
        )

    cmd = [
        "ffmpeg",
        "-y",
        "-ss", f"{window_start:.3f}",
        "-t", f"{duration:.3f}",
        "-i", video_path,
        "-vn",
        "-ac", "1",
        "-ar", str(sample_rate),
        "-loglevel", "error",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    if result.returncode != 0 or not os.path.exists(output_path):
        raise FrameExtractionError(
            f"ffmpeg audio extract failed at {window_start}s: {result.stderr.strip()}"
        )
    return output_path