TwitchDownloader/modules/processor.py

"""
Video/audio processing functionality using FFmpeg.
"""

import os
import subprocess
from colorama import Fore, Style

from .utils import detect_hardware_acceleration, get_hwaccel_encoder


class StreamProcessor:
    """Handles video/audio processing using FFmpeg."""
    
    def __init__(self, os_type: str, ffmpeg_path: str, config: dict):
        """
        Initialize the stream processor.
        
        Args:
            os_type: Operating system type ('windows' or 'linux')
            ffmpeg_path: Path to FFmpeg executable
            config: Configuration dictionary with FFmpeg settings
        """
        self.os_type = os_type
        self.ffmpeg_path = ffmpeg_path
        self.quality = config.get('quality', 'best')
        self.only_raw = config.get('onlyRaw', False)
        self.ffmpeg_threads = config.get('ffmpeg_threads', 0)
        self.ffmpeg_audio_codec = config.get('ffmpeg_audio_codec', 'aac')
        self.ffmpeg_audio_samplerate = config.get('ffmpeg_audio_samplerate', 48000)
        self.ffmpeg_audio_bitrate = config.get('ffmpeg_audio_bitrate', '192k')
        self.ffmpeg_error_recovery = config.get('ffmpeg_error_recovery', True)
        self.ffmpeg_faststart = config.get('ffmpeg_faststart', True)
        self.ffmpeg_progress = config.get('ffmpeg_progress', False)
        self.hwaccel_type = detect_hardware_acceleration(
            config.get('ffmpeg_hwaccel', 'auto'), 
            os_type
        )
    
    def process_raw_stream(self, raw_path: str, output_path: str) -> None:
        """
        Process raw .ts file into mp4/mp3 using ffmpeg.
        
        Args:
            raw_path: Path to the raw .ts file
            output_path: Path for the processed output file
        """
        if not os.path.exists(raw_path):
            print(f'{Fore.YELLOW}⚠ Raw file not found, skipping processing{Style.RESET_ALL}')
            return
        
        if self.only_raw:
            print(f'{Fore.CYAN}Keeping raw .ts file (onlyRaw mode){Style.RESET_ALL}')
            return
        
        print(f'{Fore.YELLOW}Processing raw stream file...{Style.RESET_ALL}')
        
        # Build ffmpeg command based on quality
        if self.quality == 'audio_only':
            self._process_audio(raw_path, output_path)
        else:
            self._process_video(raw_path, output_path)
        
        print(f'{Fore.GREEN}✓ Stream processed successfully{Style.RESET_ALL}')
    
    def _process_audio(self, raw_path: str, output_path: str) -> None:
        """Process audio-only stream."""
        # Audio-only conversion with modern AAC encoding
        cmd = [
            self.ffmpeg_path,
            '-i', raw_path,
            '-vn',  # No video
            '-c:a', self.ffmpeg_audio_codec,
            '-ar', str(self.ffmpeg_audio_samplerate),
            '-ac', '2',  # Stereo
            '-b:a', self.ffmpeg_audio_bitrate,
        ]
        
        # Add threading for faster encoding
        if self.ffmpeg_threads > 0:
            cmd.extend(['-threads', str(self.ffmpeg_threads)])
        
        # Add faststart for better streaming compatibility
        if self.ffmpeg_faststart and output_path.endswith(('.mp4', '.m4a')):
            cmd.extend(['-movflags', '+faststart'])
        
        cmd.append(output_path)
        
        # Run FFmpeg
        if self.ffmpeg_progress:
            subprocess.call(cmd)
        else:
            subprocess.call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
    
    def _process_video(self, raw_path: str, output_path: str) -> None:
        """Process video stream."""
        cmd = [
            self.ffmpeg_path,
            '-y',  # Overwrite output file
        ]
        
        # Add hardware acceleration if enabled
        if self.hwaccel_type and self.hwaccel_type != 'none':
            print(f'{Fore.CYAN}Using hardware acceleration: {self.hwaccel_type}{Style.RESET_ALL}')
            cmd.extend(['-hwaccel', 'auto'])
        
        cmd.extend([
            '-i', raw_path,
            '-analyzeduration', '2147483647',
            '-probesize', '2147483647',
        ])
        
        # Threading support
        if self.ffmpeg_threads >= 0:
            cmd.extend(['-threads', str(self.ffmpeg_threads)])
        
        # Error recovery options for corrupted streams
        if self.ffmpeg_error_recovery:
            cmd.extend([
                '-fflags', '+genpts',
                '-avoid_negative_ts', 'make_zero',
                '-err_detect', 'ignore_err'
            ])
        
        # Stream copy (fast, no re-encoding)
        cmd.extend([
            '-c:v', 'copy',
            '-c:a', 'copy',
            '-start_at_zero',
            '-copyts',
        ])
        
        # Add faststart for MP4 files
        if self.ffmpeg_faststart and output_path.endswith('.mp4'):
            cmd.extend(['-movflags', '+faststart'])
        
        cmd.append(output_path)
        
        # Run FFmpeg
        if self.ffmpeg_progress:
            subprocess.call(cmd)
        else:
            subprocess.call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
    
    def build_chat_output_args(self) -> str:
        """
        Build FFmpeg output arguments for chat rendering with hardware acceleration support.
        
        Returns:
            str: Complete FFmpeg output arguments string with "{save_path}" placeholder
        """
        if self.hwaccel_type and self.hwaccel_type != 'none':
            encoder = get_hwaccel_encoder(self.hwaccel_type)
            print(f'{Fore.CYAN}Using hardware-accelerated encoder for chat: {encoder}{Style.RESET_ALL}')
            
            if 'nvenc' in encoder:
                result = f'-c:v {encoder} -preset p4 -cq 18 -pix_fmt yuv420p "{{save_path}}"'
            elif 'qsv' in encoder:
                result = f'-c:v {encoder} -global_quality 18 -pix_fmt yuv420p "{{save_path}}"'
            elif 'amf' in encoder:
                result = f'-c:v {encoder} -qp_i 18 -pix_fmt yuv420p "{{save_path}}"'
            elif 'vaapi' in encoder:
                result = f'-c:v {encoder} -qp 18 -pix_fmt yuv420p "{{save_path}}"'
            else:
                result = f'-c:v libx264 -preset veryfast -crf 18 -pix_fmt yuv420p "{{save_path}}"'
        else:
            # Default software encoding
            result = f'-c:v libx264 -preset veryfast -crf 18 -pix_fmt yuv420p "{{save_path}}"'
        
        print(f'{Fore.CYAN}DEBUG - Generated output_args: {result}{Style.RESET_ALL}')
        return result