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) -> bool:
        """
        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

        Returns:
            bool: True when conversion succeeded, False otherwise
        """
        if not os.path.exists(raw_path):
            print(f'{Fore.YELLOW}⚠ Raw file not found, skipping processing{Style.RESET_ALL}')
            return False
        
        if self.only_raw:
            print(f'{Fore.CYAN}Keeping raw .ts file (onlyRaw mode){Style.RESET_ALL}')
            return False
        
        print(f'{Fore.YELLOW}Processing raw stream file...{Style.RESET_ALL}')
        
        # Build ffmpeg command based on quality
        if self.quality == 'audio_only':
            result = self._process_audio(raw_path, output_path)
        else:
            result = self._process_video(raw_path, output_path)
        
        if result:
            print(f'{Fore.GREEN}✓ Stream processed successfully{Style.RESET_ALL}')
        else:
            print(f'{Fore.RED}✗ Stream processing failed{Style.RESET_ALL}')

        return result
    
    def _run_ffmpeg_command(self, cmd: list, output_path: str) -> bool:
        """Run FFmpeg while streaming its output to the terminal."""
        print(f'{Fore.CYAN}Running FFmpeg: {' '.join(cmd)}{Style.RESET_ALL}')

        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding='utf-8',
            errors='replace'
        )

        if process.stdout:
            for line in process.stdout:
                print(line, end='')

        result = process.wait()
        if result != 0:
            print(f'{Fore.RED}✗ FFmpeg exited with code: {result}{Style.RESET_ALL}')
            return False

        if not os.path.exists(output_path):
            print(f'{Fore.RED}✗ FFmpeg did not create output: {output_path}{Style.RESET_ALL}')
            return False

        if os.path.getsize(output_path) == 0:
            print(f'{Fore.RED}✗ FFmpeg created an empty output file: {output_path}{Style.RESET_ALL}')
            return False

        return True

    def _process_audio(self, raw_path: str, output_path: str) -> bool:
        """Process audio-only stream."""
        # Audio-only conversion with modern AAC encoding
        cmd = [
            self.ffmpeg_path,
            '-y',
            '-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)
        return self._run_ffmpeg_command(cmd, output_path)
    
    def _process_video(self, raw_path: str, output_path: str) -> bool:
        """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)
        return self._run_ffmpeg_command(cmd, output_path)
    
    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}}"'
        
        return result
    
    def merge_video_and_chat(self, video_path: str, chat_path: str, output_path: str, layout: str = 'side-by-side') -> bool:
        """
        Merge video and chat into a single file.
        
        Args:
            video_path: Path to the main video file
            chat_path: Path to the chat video file
            output_path: Path for the merged output file
            layout: Merge layout - 'side-by-side' or 'overlay' (default: 'side-by-side')
            
        Returns:
            bool: True if merge succeeded, False otherwise
        """
        if not os.path.exists(video_path):
            print(f'{Fore.RED}✗ Video file not found: {video_path}{Style.RESET_ALL}')
            return False
        
        if not os.path.exists(chat_path):
            print(f'{Fore.RED}✗ Chat file not found: {chat_path}{Style.RESET_ALL}')
            return False
        
        print(f'{Fore.YELLOW}Merging video and chat ({layout})...{Style.RESET_ALL}')
        print(f'{Fore.CYAN}  This may take several minutes depending on video length{Style.RESET_ALL}')
        
        try:
            if layout == 'overlay':
                # Overlay chat on top of video (right side)
                filter_complex = (
                    '[0:v]scale=-2:1080[main];'
                    '[1:v]scale=500:1080[chat];'
                    '[main][chat]overlay=main_w-overlay_w:0[outv]'
                )
            else:
                # Side-by-side layout (default)
                filter_complex = (
                    '[0:v]scale=-2:1080[main];'
                    '[1:v]scale=500:1080[chat];'
                    '[main][chat]hstack=inputs=2[outv]'
                )
            
            cmd = [
                self.ffmpeg_path,
                '-y',
                '-i', video_path,
                '-i', chat_path,
                '-filter_complex', filter_complex,
                '-map', '[outv]',
                '-map', '0:a?',  # Use audio from main video if available
            ]
            
            # Add hardware acceleration for encoding if available
            if self.hwaccel_type and self.hwaccel_type != 'none':
                encoder = get_hwaccel_encoder(self.hwaccel_type)
                cmd.extend(['-c:v', encoder])
                print(f'{Fore.CYAN}  Using hardware encoder: {encoder}{Style.RESET_ALL}')
                
                if 'nvenc' in encoder:
                    cmd.extend(['-preset', 'p4', '-cq', '18'])
                elif 'qsv' in encoder:
                    cmd.extend(['-global_quality', '18'])
                elif 'amf' in encoder:
                    cmd.extend(['-qp_i', '18'])
                else:
                    cmd.extend(['-preset', 'medium', '-crf', '18'])
            else:
                print(f'{Fore.CYAN}  Using software encoder: libx264{Style.RESET_ALL}')
                cmd.extend(['-c:v', 'libx264', '-preset', 'medium', '-crf', '18'])
            
            # Audio codec
            cmd.extend(['-c:a', 'copy'])  # Copy audio without re-encoding
            
            # Pixel format and faststart
            cmd.extend(['-pix_fmt', 'yuv420p'])
            if self.ffmpeg_faststart and output_path.endswith('.mp4'):
                cmd.extend(['-movflags', '+faststart'])
            
            cmd.append(output_path)
            
            # Run FFmpeg with visible output to show progress
            print(f'{Fore.CYAN}  Processing...{Style.RESET_ALL}')
            result = subprocess.call(cmd)
            
            if result == 0:
                print(f'{Fore.GREEN}✓ Video and chat merged successfully{Style.RESET_ALL}')
                return True
            else:
                print(f'{Fore.RED}✗ Merge failed with exit code: {result}{Style.RESET_ALL}')
                print(f'{Fore.YELLOW}  Check FFmpeg output above for details{Style.RESET_ALL}')
                return False
                
        except Exception as e:
            print(f'{Fore.RED}✗ Merge failed: {str(e)}{Style.RESET_ALL}')
            import traceback
            traceback.print_exc()
            return False