src/video_preprocessing.py

"""Video preprocessing: metadata, sampled frame extraction and audio export.

Design goals (MacBook Air M4 friendly):
    * Sample frames at a configurable interval — never decode every frame.
    * Extract audio to a 16 kHz mono WAV (ideal for Whisper-family ASR).
    * Keep memory flat: frames are written to disk and yielded as paths, not
      held in a giant in-memory list.

Heavy dependencies (`opencv-python`, `ffmpeg`) are imported lazily so that the
dashboard and sample-output generator can import this module without them.
"""

from __future__ import annotations

import shutil
import subprocess
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Dict, List, Optional

from src.config import Config, CONFIG
from src.storage import work_dir
from src.utils import format_timestamp


@dataclass
class VideoMetadata:
    video_id: str
    path: str
    duration_sec: float
    fps: float
    width: int
    height: int
    frame_count: int

    def as_dict(self) -> Dict[str, object]:
        return asdict(self)


@dataclass
class FrameSample:
    index: int
    time_sec: float
    time_label: str
    path: str


def _require_cv2():
    try:
        import cv2  # type: ignore

        return cv2
    except ImportError as exc:  # pragma: no cover - environment dependent
        raise RuntimeError(
            "opencv-python is required for live video preprocessing. "
            "Install with `pip install -r requirements-local.txt`."
        ) from exc


def probe_metadata(video_path: Path, video_id: str) -> VideoMetadata:
    """Read duration / fps / resolution using OpenCV (no ffprobe needed)."""
    cv2 = _require_cv2()
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video: {video_path}")
    try:
        fps = float(cap.get(cv2.CAP_PROP_FPS)) or 0.0
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        duration = (frame_count / fps) if fps else 0.0
    finally:
        cap.release()
    return VideoMetadata(
        video_id=video_id,
        path=str(video_path),
        duration_sec=round(duration, 3),
        fps=round(fps, 3),
        width=width,
        height=height,
        frame_count=frame_count,
    )


def extract_frames(
    video_path: Path,
    video_id: str,
    interval_sec: float,
    max_duration_sec: Optional[int] = None,
) -> List[FrameSample]:
    """Extract one frame every ``interval_sec`` seconds, written as JPEGs.

    Returns a list of :class:`FrameSample`. We seek by timestamp instead of
    decoding sequentially, which keeps CPU usage low on short videos.
    """
    cv2 = _require_cv2()
    out_dir = work_dir(video_id) / "frames"
    out_dir.mkdir(parents=True, exist_ok=True)

    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video: {video_path}")

    samples: List[FrameSample] = []
    try:
        fps = float(cap.get(cv2.CAP_PROP_FPS)) or 0.0
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
        duration = (frame_count / fps) if fps else 0.0
        if max_duration_sec:
            duration = min(duration, float(max_duration_sec))

        t = 0.0
        idx = 0
        while t <= duration:
            cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000.0)
            ok, frame = cap.read()
            if not ok:
                break
            fname = out_dir / f"frame_{idx:04d}.jpg"
            cv2.imwrite(str(fname), frame)
            samples.append(
                FrameSample(
                    index=idx,
                    time_sec=round(t, 3),
                    time_label=format_timestamp(t),
                    path=str(fname),
                )
            )
            idx += 1
            t += interval_sec
    finally:
        cap.release()
    return samples


def extract_audio(video_path: Path, video_id: str) -> Optional[Path]:
    """Export 16 kHz mono WAV via ffmpeg. Returns ``None`` if ffmpeg missing."""
    if shutil.which("ffmpeg") is None:
        return None
    out_path = work_dir(video_id) / "audio.wav"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(video_path),
        "-vn",
        "-ac",
        "1",
        "-ar",
        "16000",
        "-f",
        "wav",
        str(out_path),
    ]
    proc = subprocess.run(cmd, capture_output=True)
    if proc.returncode != 0 or not out_path.exists():
        return None
    return out_path


def cleanup_workdir(video_id: str) -> None:
    """Remove transient frames/audio for a video id."""
    d = work_dir(video_id)
    if d.exists():
        shutil.rmtree(d, ignore_errors=True)


def preprocess(
    video_path: Path,
    video_id: str,
    config: Config = CONFIG,
) -> Dict[str, object]:
    """Run the full preprocessing step and return metadata + frames + audio."""
    metadata = probe_metadata(video_path, video_id)
    frames = extract_frames(
        video_path,
        video_id,
        interval_sec=config.frame_interval_sec,
        max_duration_sec=config.max_video_duration_sec,
    )
    audio_path = extract_audio(video_path, video_id)
    return {
        "metadata": metadata,
        "frames": frames,
        "audio_path": audio_path,
    }