VideoDetect/src/prober.py

"""
Video Probing Module

Extracts video metadata (codec, resolution, duration, frame rate) using ffprobe.
Handles errors gracefully for corrupt or unsupported files.
"""

import json
import logging
import subprocess
import time
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger(__name__)


@dataclass
class VideoMetadata:
    """Metadata extracted from a video file."""
    file_path: str
    codec: Optional[str] = None
    resolution_w: Optional[int] = None
    resolution_h: Optional[int] = None
    duration: Optional[float] = None
    frame_rate: Optional[float] = None
    file_size: Optional[int] = None
    status: str = "PENDING"
    error_message: Optional[str] = None
    probe_time: float = field(default_factory=time.time)

    @property
    def is_valid(self) -> bool:
        return self.status == "PENDING" and self.codec is not None

    @property
    def is_unscannable(self) -> bool:
        return self.status == "UNSCANNABLE"

    def to_dict(self) -> dict:
        return {
            "file_path": self.file_path,
            "codec": self.codec,
            "resolution_w": self.resolution_w,
            "resolution_h": self.resolution_h,
            "duration": self.duration,
            "frame_rate": self.frame_rate,
            "file_size": self.file_size,
            "status": self.status,
            "error_message": self.error_message,
        }


class VideoProber:
    """Probe video files using ffprobe to extract metadata."""

    def __init__(self, timeout: int = 10):
        self.timeout = timeout

    def probe(self, file_path: str) -> VideoMetadata:
        """Probe a video file and return metadata."""
        metadata = VideoMetadata(file_path=file_path)

        # Get file size
        try:
            metadata.file_size = self._get_file_size(file_path)
        except OSError as e:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = f"Cannot read file: {e}"
            logger.warning("File access error for %s: %s", file_path, e)
            return metadata

        # Run ffprobe
        try:
            result = self._run_ffprobe(file_path)
            if result is None:
                metadata.status = "UNSCANNABLE"
                metadata.error_message = "ffprobe timed out or failed"
                logger.warning("ffprobe failed for %s", file_path)
                return metadata

            metadata = self._parse_ffprobe_output(result, file_path)
        except subprocess.TimeoutExpired:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = f"ffprobe timed out after {self.timeout}s"
            logger.warning("ffprobe timeout for %s", file_path)
        except Exception as e:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = f"ffprobe error: {e}"
            logger.warning("ffprobe error for %s: %s", file_path, e)

        return metadata

    def _run_ffprobe(self, file_path: str) -> Optional[str]:
        """Run ffprobe and return JSON output."""
        cmd = [
            "ffprobe",
            "-v", "error",
            "-select_streams", "v:0",
            "-show_entries", "stream=codec_name,width,height,r_frame_rate,duration",
            "-show_entries", "format=duration,size",
            "-of", "json",
            file_path,
        ]

        try:
            output = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=self.timeout,
            )

            if output.returncode != 0:
                logger.debug("ffprobe stderr for %s: %s", file_path, output.stderr)
                return None

            return output.stdout
        except subprocess.TimeoutExpired:
            return None

    def _parse_ffprobe_output(self, output: str, file_path: str) -> VideoMetadata:
        """Parse ffprobe JSON output into VideoMetadata."""
        metadata = VideoMetadata(file_path=file_path)

        try:
            data = json.loads(output)
        except json.JSONDecodeError as e:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = f"Invalid ffprobe output: {e}"
            return metadata

        # Extract stream info
        streams = data.get("streams", [])
        if not streams:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = "No video streams found"
            return metadata

        stream = streams[0]
        metadata.codec = stream.get("codec_name")
        metadata.resolution_w = stream.get("width")
        metadata.resolution_h = stream.get("height")

        # Parse frame rate
        r_frame_rate = stream.get("r_frame_rate")
        if r_frame_rate and "/" in r_frame_rate:
            try:
                num, den = r_frame_rate.split("/")
                metadata.frame_rate = float(num) / float(den) if float(den) != 0 else None
            except (ValueError, ZeroDivisionError):
                pass

        # Extract format info
        format_info = data.get("format", {})
        if format_info.get("duration"):
            try:
                metadata.duration = float(format_info["duration"])
            except (ValueError, TypeError):
                pass

        if format_info.get("size"):
            try:
                metadata.file_size = int(format_info["size"])
            except (ValueError, TypeError):
                pass

        # Validate we got essential data
        if not metadata.codec:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = "No codec information available"
        elif not metadata.duration:
            metadata.status = "UNSCANNABLE"
            metadata.error_message = "No duration information available"

        return metadata

    @staticmethod
    def _get_file_size(file_path: str) -> int:
        """Get file size in bytes."""
        return os.path.getsize(file_path)