utils/video_processor.py

"""
Video Processor - Handles video I/O using FFmpeg
Supports frame extraction, video assembly, and audio preservation
"""
import os
import cv2
import json
import shutil
import subprocess
import numpy as np
from pathlib import Path
from typing import Generator, Tuple, Optional, Dict
from dataclasses import dataclass


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


class VideoProcessor:
    """Handles all video I/O operations"""
    
    def __init__(self, config):
        self.config = config
        self._verify_ffmpeg()
    
    def _verify_ffmpeg(self):
        """Verify FFmpeg is installed"""
        try:
            result = subprocess.run(
                ["ffmpeg", "-version"], 
                capture_output=True, text=True, timeout=5
            )
            if result.returncode != 0:
                raise RuntimeError("FFmpeg not working properly")
        except FileNotFoundError:
            raise RuntimeError(
                "FFmpeg not found! Install it:\n"
                "  Ubuntu: sudo apt install ffmpeg\n"
                "  Mac: brew install ffmpeg\n"
                "  Windows: choco install ffmpeg"
            )
    
    def get_video_info(self, video_path: str) -> VideoInfo:
        """Extract video metadata using ffprobe"""
        cmd = [
            "ffprobe", "-v", "quiet",
            "-print_format", "json",
            "-show_format", "-show_streams",
            video_path
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            raise ValueError(f"Cannot read video: {video_path}")
        
        probe = json.loads(result.stdout)
        
        # Find video stream
        video_stream = None
        has_audio = False
        for stream in probe.get("streams", []):
            if stream["codec_type"] == "video" and video_stream is None:
                video_stream = stream
            elif stream["codec_type"] == "audio":
                has_audio = True
        
        if not video_stream:
            raise ValueError("No video stream found")
        
        # Parse FPS
        fps_parts = video_stream.get("r_frame_rate", "30/1").split("/")
        fps = float(fps_parts[0]) / float(fps_parts[1]) if len(fps_parts) == 2 else 30.0
        
        # Parse frame count
        nb_frames = int(video_stream.get("nb_frames", 0))
        duration = float(probe.get("format", {}).get("duration", 0))
        
        if nb_frames == 0 and duration > 0:
            nb_frames = int(duration * fps)
        
        return VideoInfo(
            width=int(video_stream["width"]),
            height=int(video_stream["height"]),
            fps=fps,
            total_frames=nb_frames,
            duration=duration,
            has_audio=has_audio,
            codec=video_stream.get("codec_name", "unknown"),
            filepath=video_path
        )
    
    def extract_frames_to_dir(
        self, 
        video_path: str, 
        output_dir: str,
        max_height: Optional[int] = None,
        progress_callback=None
    ) -> Tuple[str, VideoInfo]:
        """
        Extract all frames as numbered JPEGs (required by SAM 2)
        
        Args:
            video_path: Path to input video
            output_dir: Directory to save frames
            max_height: Optional max height for resizing
            progress_callback: Optional callback(current, total)
        
        Returns:
            (frames_dir, video_info)
        """
        os.makedirs(output_dir, exist_ok=True)
        info = self.get_video_info(video_path)
        
        # Build FFmpeg command
        cmd = ["ffmpeg", "-y", "-i", video_path]
        
        # Add scaling if needed
        if max_height and info.height > max_height:
            cmd.extend(["-vf", f"scale=-2:{max_height}"])
        
        # Output as numbered JPEGs
        cmd.extend([
            "-qscale:v", "2",  # High quality JPEG
            os.path.join(output_dir, "%06d.jpg")
        ])
        
        process = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        _, stderr = process.communicate()
        
        if process.returncode != 0:
            raise RuntimeError(f"Frame extraction failed: {stderr.decode()}")
        
        # Count extracted frames
        frame_files = sorted([
            f for f in os.listdir(output_dir) 
            if f.endswith('.jpg')
        ])
        info.total_frames = len(frame_files)
        
        # Update dimensions if resized
        if frame_files:
            sample = cv2.imread(os.path.join(output_dir, frame_files[0]))
            if sample is not None:
                info.height, info.width = sample.shape[:2]
        
        print(f"📹 Extracted {info.total_frames} frames ({info.width}x{info.height} @ {info.fps:.1f} FPS)")
        return output_dir, info
    
    def read_frames_generator(
        self, 
        video_path: str,
        max_height: Optional[int] = None
    ) -> Generator[Tuple[int, np.ndarray], None, None]:
        """
        Stream frames from video using OpenCV (memory efficient)
        
        Yields:
            (frame_index, frame_bgr)
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video: {video_path}")
        
        frame_idx = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Resize if needed
            if max_height and frame.shape[0] > max_height:
                scale = max_height / frame.shape[0]
                new_w = int(frame.shape[1] * scale)
                frame = cv2.resize(frame, (new_w, max_height))
            
            yield frame_idx, frame
            frame_idx += 1
        
        cap.release()
    
    def assemble_video(
        self, 
        frames_dir: str,
        output_path: str,
        fps: float,
        original_video: Optional[str] = None,
        progress_callback=None
    ) -> str:
        """
        Assemble processed frames back into video with optional audio
        
        Args:
            frames_dir: Directory with numbered JPEG frames
            output_path: Output video path
            fps: Frame rate
            original_video: Original video to copy audio from
            progress_callback: Optional callback
        
        Returns:
            Path to output video
        """
        temp_video = output_path + ".temp.mp4"
        
        # Step 1: Encode frames to video
        cmd = [
            "ffmpeg", "-y",
            "-framerate", str(fps),
            "-i", os.path.join(frames_dir, "%06d.jpg"),
            "-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",
        ]
        
        if original_video:
            cmd.append(temp_video)
        else:
            cmd.append(output_path)
        
        process = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        _, stderr = process.communicate()
        
        if process.returncode != 0:
            raise RuntimeError(f"Video encoding failed: {stderr.decode()}")
        
        # Step 2: Mux audio from original video
        if original_video:
            info = self.get_video_info(original_video)
            
            if info.has_audio:
                mux_cmd = [
                    "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
                ]
                
                process = subprocess.Popen(
                    mux_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                _, stderr = process.communicate()
                
                if process.returncode != 0:
                    # Fallback: just use video without audio
                    shutil.move(temp_video, output_path)
                    print("⚠️  Audio muxing failed, output has no audio")
                else:
                    os.remove(temp_video)
                    print("🔊 Audio preserved from original video")
            else:
                shutil.move(temp_video, output_path)
        
        file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
        print(f"✅ Output video: {output_path} ({file_size_mb:.1f} MB)")
        
        return output_path
    
    def cleanup_temp(self, temp_dir: str):
        """Remove temporary files"""
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)