encoding_config.py

"""
Encoding configuration module for video processing.
Manages encoding methods: bitrate-based (VBR) and quality-based (CRF) encoding.
"""

from enum import Enum
from dataclasses import dataclass
from typing import Dict, Optional, Any
import ffmpeg


class EncodingMethod(Enum):
    """Encoding method types"""
    VBR = "vbr"  # Variable Bitrate (bitrate-based)
    CRF = "crf"  # Constant Rate Factor (quality-based)


class VideoCodec(Enum):
    """Video codec types"""
    H264 = "h264"  # H.264/AVC
    H265 = "h265"  # H.265/HEVC
    AV1 = "av1"    # AV1


@dataclass
class EncodingConfig:
    """Encoding configuration container"""
    method: EncodingMethod
    value: float
    video_codec: str = "libx265"
    codec_type: VideoCodec = VideoCodec.H265
    hw_accel: Optional[str] = None
    preset: str = "medium"
    additional_params: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.additional_params is None:
            self.additional_params = {}


class EncodingConfigManager:
    """Manages encoding configuration and FFmpeg parameter generation"""
    
    # Quality presets for CRF encoding (H.264/H.265: 0-51 range)
    CRF_PRESETS = {
        'ultra_high': 18,   # Near lossless
        'high': 23,         # High quality
        'medium': 28,       # Balanced quality/size
        'low': 33,          # Lower quality, smaller size
        'very_low': 38      # Very low quality
    }
    
    # Quality presets for AV1 CRF encoding (0-63 range)
    AV1_CRF_PRESETS = {
        'ultra_high': 20,   # Near lossless
        'high': 26,         # High quality
        'medium': 32,       # Preserve detail with substantial compression
        'low': 38,          # Lower quality, smaller size
        'very_low': 46      # Very low quality
    }
    
    # Bitrate multiplier presets for VBR encoding
    VBR_PRESETS = {
        'highest': 1.2,     # 120% of original bitrate
        'high': 1.0,        # 100% of original bitrate
        'medium': 0.75,     # 75% of original bitrate
        'low': 0.5,         # 50% of original bitrate
        'lowest': 0.25      # 25% of original bitrate
    }
    
    def __init__(self):
        self.config = EncodingConfig(
            method=EncodingMethod.CRF,
            value=23,  # Default CRF value
            video_codec="libx265",
            codec_type=VideoCodec.H265
        )
    
    def set_crf_encoding(self, crf_value: float, preset: str = "medium"):
        """Configure for CRF (quality-based) encoding"""
        # Validate CRF range: 0-51 for x264/x265, 0-63 for AV1
        if self.config.codec_type == VideoCodec.AV1:
            crf_value = max(0, min(63, crf_value))
        else:
            crf_value = max(0, min(51, crf_value))
        
        self.config.method = EncodingMethod.CRF
        self.config.value = crf_value
        self.config.preset = preset
    
    def set_vbr_encoding(self, bitrate_multiplier: float, preset: str = "medium"):
        """Configure for VBR (bitrate-based) encoding"""
        # Validate multiplier range
        bitrate_multiplier = max(0.1, min(10.0, bitrate_multiplier))
        
        self.config.method = EncodingMethod.VBR
        self.config.value = bitrate_multiplier
        self.config.preset = preset
    
    def set_hardware_acceleration(self, hw_accel: str, video_codec: str):
        """Set hardware acceleration settings"""
        self.config.hw_accel = hw_accel
        self.config.video_codec = video_codec
    
    def set_codec_type(self, codec_type: VideoCodec, hw_accel: Optional[str] = None):
        """Set video codec type (H.264, H.265, or AV1)"""
        self.config.codec_type = codec_type
        
        # Update video codec based on hardware acceleration and codec type
        if hw_accel:
            if codec_type == VideoCodec.H264:
                if 'nvenc' in self.config.video_codec:
                    self.config.video_codec = 'h264_nvenc'
                elif 'amf' in self.config.video_codec:
                    self.config.video_codec = 'h264_amf'
                elif 'qsv' in self.config.video_codec:
                    self.config.video_codec = 'h264_qsv'
                else:
                    self.config.video_codec = 'libx264'
            elif codec_type == VideoCodec.H265:
                if 'nvenc' in self.config.video_codec:
                    self.config.video_codec = 'hevc_nvenc'
                elif 'amf' in self.config.video_codec:
                    self.config.video_codec = 'hevc_amf'
                elif 'qsv' in self.config.video_codec:
                    self.config.video_codec = 'hevc_qsv'
                else:
                    self.config.video_codec = 'libx265'
            else:  # AV1
                if 'nvenc' in self.config.video_codec:
                    self.config.video_codec = 'av1_nvenc'
                elif 'amf' in self.config.video_codec:
                    self.config.video_codec = 'av1_amf'
                elif 'qsv' in self.config.video_codec:
                    self.config.video_codec = 'av1_qsv'
                else:
                    self.config.video_codec = 'libsvtav1'
        else:
            # Software encoding
            if codec_type == VideoCodec.H264:
                self.config.video_codec = 'libx264'
            elif codec_type == VideoCodec.H265:
                self.config.video_codec = 'libx265'
            else:  # AV1
                self.config.video_codec = 'libsvtav1'
    
    def get_crf_from_preset(self, preset_name: str) -> float:
        """Get CRF value from preset name (uses AV1 presets if codec is AV1)"""
        if self.config.codec_type == VideoCodec.AV1:
            return self.AV1_CRF_PRESETS.get(preset_name.lower(), 38)
        return self.CRF_PRESETS.get(preset_name.lower(), 23)
    
    def get_vbr_from_preset(self, preset_name: str) -> float:
        """Get VBR multiplier from preset name"""
        return self.VBR_PRESETS.get(preset_name.lower(), 0.75)
    
    def calculate_target_bitrate(self, original_bitrate: str) -> int:
        """Calculate target bitrate for VBR encoding"""
        if self.config.method != EncodingMethod.VBR:
            raise ValueError("Can only calculate target bitrate for VBR encoding")
        
        try:
            original_bitrate_int = int(original_bitrate)
            return int(original_bitrate_int * self.config.value)
        except (ValueError, TypeError):
            raise ValueError(f"Invalid original bitrate: {original_bitrate}")
    
    def generate_ffmpeg_params(self, input_file: str, output_file: str, 
                             original_bitrate: Optional[str] = None,
                             scale_filter: Optional[str] = None) -> Dict[str, Any]:
        """Generate FFmpeg parameters based on current configuration"""
        # Base input configuration with increased thread queue size for slow storage (HDD)
        input_config = ffmpeg.input(input_file, thread_queue_size=1024)
        
        # Hardware acceleration
        global_args = []
        if self.config.hw_accel:
            global_args.extend(['-hwaccel', self.config.hw_accel])
        
        # Video encoding parameters
        video_params = {
            'vcodec': self.config.video_codec
        }
        
        # Add preset only for software encoders and some hardware encoders
        if ('amf' not in self.config.video_codec and 
            'qsv' not in self.config.video_codec):
            if 'svtav1' in self.config.video_codec:
                # SVT-AV1 uses numeric preset 0-13 (0=best quality, 13=fastest)
                video_params['preset'] = self._map_preset_to_svtav1(self.config.preset)
            elif 'aom' in self.config.video_codec:
                # libaom-av1 uses cpu-used 0-8 (0=best quality, 8=fastest)
                video_params['cpu-used'] = self._map_preset_to_aom(self.config.preset)
            else:
                video_params['preset'] = self.config.preset
        
        # Method-specific parameters
        if self.config.method == EncodingMethod.CRF:
            # Hardware encoders use their own quality controls, configured in
            # their codec-specific optimization methods below.
            if not any(
                hardware_encoder in self.config.video_codec
                for hardware_encoder in ('amf', 'nvenc', 'qsv')
            ):
                video_params['crf'] = int(self.config.value)
        elif self.config.method == EncodingMethod.VBR:
            if original_bitrate:
                target_bitrate = self.calculate_target_bitrate(original_bitrate)
                video_params['b:v'] = f"{target_bitrate}"
                if 'amf' in self.config.video_codec:
                    # Keep AMF's peak rate close to the requested target. The
                    # buffer allows normal short-term bitrate fluctuations.
                    video_params['maxrate'] = f"{int(target_bitrate * 1.10)}"
                    video_params['bufsize'] = f"{int(target_bitrate * 2)}"
            else:
                raise ValueError("Original bitrate required for VBR encoding")
        
        # Add scale filter if needed
        filter_chain = []
        if scale_filter:
            filter_chain.append(scale_filter)
        
        # Combine filters
        if filter_chain:
            video_params['vf'] = ','.join(filter_chain)
        
        # Additional codec-specific optimizations
        if 'nvenc' in self.config.video_codec:
            video_params.update(self._get_nvenc_optimizations())
        elif 'amf' in self.config.video_codec:
            video_params.update(self._get_amf_optimizations())
        elif 'qsv' in self.config.video_codec:
            video_params.update(self._get_qsv_optimizations())
        elif 'x265' in self.config.video_codec or 'libx265' in self.config.video_codec:
            video_params.update(self._get_x265_optimizations())
        elif 'svtav1' in self.config.video_codec or 'aom' in self.config.video_codec:
            video_params.update(self._get_av1_optimizations())
        
        # Add additional parameters
        video_params.update(self.config.additional_params)
        
        return {
            'input_config': input_config,
            'output_file': output_file,
            'video_params': video_params,
            'global_args': global_args
        }
    
    def _get_nvenc_optimizations(self) -> Dict[str, Any]:
        """Get NVIDIA NVENC specific optimizations"""
        if self.config.codec_type == VideoCodec.AV1:
            optimizations = {
                'profile:v': 'main',
                'preset': 'p7',
                'tune': 'uhq',
                'multipass': 'fullres',
                'rc-lookahead': '32',
                'b_ref_mode': 'middle',
                'spatial_aq': '1',
                'temporal_aq': '1',
                'aq-strength': '8',
            }
            if self.config.method == EncodingMethod.CRF:
                optimizations.update({
                    'rc': 'vbr',
                    'cq': int(self.config.value),
                })
            else:
                optimizations['rc'] = 'vbr'
            return optimizations

        return {
            'rc': 'vbr' if self.config.method == EncodingMethod.VBR else 'cqp',
            'profile:v': 'main',
            'level': '4.1',
            'b_ref_mode': 'middle',
            'spatial_aq': '1',
            'temporal_aq': '1'
        }
    
    def _get_amf_optimizations(self) -> Dict[str, Any]:
        """Get AMD AMF specific optimizations"""
        if self.config.codec_type == VideoCodec.AV1:
            optimizations = {
                'profile:v': 'main',
                'usage': 'high_quality',
                'quality': 'high_quality',
                'preanalysis': '1',
                'aq_mode': 'caq',
                'max_b_frames': '3',
                'high_motion_quality_boost_enable': '1',
            }
            if self.config.method == EncodingMethod.CRF:
                optimizations.update({
                    'rc': 'qvbr',
                    'qvbr_quality_level': self._get_amf_qvbr_quality_level(),
                })
            else:
                optimizations['rc'] = 'hqvbr'
            return optimizations

        optimizations = {
            'profile:v': 'main',
        }
        
        # Rate control mode
        if self.config.method == EncodingMethod.VBR:
            optimizations['rc'] = 'vbr_peak'
        else:
            optimizations['rc'] = 'cqp'
            qp = int(self.config.value)
            optimizations.update({
                'qp_i': qp,
                'qp_p': qp,
                'qp_b': qp,
            })
        
        # Quality settings
        optimizations['quality'] = 'balanced'
        
        return optimizations
    
    def _get_qsv_optimizations(self) -> Dict[str, Any]:
        """Get Intel QuickSync specific optimizations"""
        if self.config.codec_type == VideoCodec.AV1:
            optimizations = {
                'profile:v': 'main',
                'preset': 'veryslow',
                'look_ahead_depth': '40',
                'extbrc': '1',
                'adaptive_i': '1',
                'adaptive_b': '1',
            }
            if self.config.method == EncodingMethod.CRF:
                optimizations['global_quality'] = int(self.config.value)
            return optimizations

        return {
            'profile:v': 'main',
            'level': '4.1',
            'look_ahead': '1',
            'look_ahead_depth': '40'
        }
    
    def _get_x265_optimizations(self) -> Dict[str, Any]:
        """Get x265 software encoder specific optimizations"""
        return {
            'profile:v': 'main',
            'level': '4.1',
            'x265-params': 'aq-mode=3:aq-strength=0.8:deblock=1,1'
        }
    
    def _get_av1_optimizations(self) -> Dict[str, Any]:
        """Get AV1 software encoder specific optimizations"""
        optimizations = {
            'profile:v': 'main',
        }
        
        if 'svtav1' in self.config.video_codec:
            # Preserve source detail rather than adding synthetic film grain.
            optimizations['svtav1-params'] = (
                'tune=0:film-grain=0:enable-overlays=1:enable-tf=1'
            )
        elif 'aom' in self.config.video_codec:
            # libaom-av1 specific params
            optimizations['aom-params'] = (
                'aq-mode=1:enable-cdef=1:enable-restoration=1'
            )
        
        return optimizations

    def _get_amf_qvbr_quality_level(self) -> int:
        """Map the CRF-style AV1 scale to AMF's higher-is-better QVBR scale."""
        return max(1, min(51, 63 - int(self.config.value)))
    
    @staticmethod
    def _map_preset_to_svtav1(preset_name: str) -> int:
        """Map human-readable preset name to SVT-AV1 numeric preset (0-13).
        0 = best quality/slowest, 13 = fastest/lowest quality."""
        mapping = {
            'ultrafast': 13,
            'superfast': 11,
            'veryfast': 9,
            'faster': 7,
            'fast': 5,
            'medium': 6,
            'slow': 4,
            'slower': 2,
            'veryslow': 0,
            'placebo': 0,
        }
        return mapping.get(preset_name.lower(), 6)
    
    @staticmethod
    def _map_preset_to_aom(preset_name: str) -> int:
        """Map human-readable preset name to libaom-av1 cpu-used (0-8).
        0 = best quality/slowest, 8 = fastest/lowest quality."""
        mapping = {
            'ultrafast': 8,
            'superfast': 7,
            'veryfast': 6,
            'faster': 5,
            'fast': 4,
            'medium': 3,
            'slow': 2,
            'slower': 1,
            'veryslow': 0,
            'placebo': 0,
        }
        return mapping.get(preset_name.lower(), 3)
    
    def get_config_summary(self) -> Dict[str, Any]:
        """Get a summary of current encoding configuration"""
        codec_names = {
            VideoCodec.H264: 'H.264/AVC',
            VideoCodec.H265: 'H.265/HEVC',
            VideoCodec.AV1: 'AV1',
        }
        
        summary = {
            'method': self.config.method.value,
            'value': self.config.value,
            'video_codec': self.config.video_codec,
            'codec_type': self.config.codec_type.value,
            'codec_name': codec_names.get(self.config.codec_type, 'Unknown'),
            'hw_accel': self.config.hw_accel,
            'preset': self.config.preset
        }
        
        if self.config.method == EncodingMethod.CRF:
            quality_mode = "QP" if 'amf' in self.config.video_codec else "CRF"
            summary['description'] = f"Quality-based encoding ({quality_mode} {self.config.value})"
            summary['quality_preset'] = self._get_crf_quality_description()
        else:
            summary['description'] = f"Bitrate-based encoding ({self.config.value}x multiplier)"
            summary['bitrate_preset'] = self._get_vbr_quality_description()
        
        return summary
    
    def _get_crf_quality_description(self) -> str:
        """Get quality description for CRF value"""
        crf = self.config.value
        if self.config.codec_type == VideoCodec.AV1:
            if crf <= 22:
                return "Ultra High Quality (Near Lossless)"
            elif crf <= 30:
                return "High Quality"
            elif crf <= 38:
                return "Medium Quality (Balanced)"
            elif crf <= 46:
                return "Low Quality"
            else:
                return "Very Low Quality"
        else:
            if crf <= 18:
                return "Ultra High Quality (Near Lossless)"
            elif crf <= 23:
                return "High Quality"
            elif crf <= 28:
                return "Medium Quality (Balanced)"
            elif crf <= 33:
                return "Low Quality"
            else:
                return "Very Low Quality"
    
    def _get_vbr_quality_description(self) -> str:
        """Get quality description for VBR multiplier"""
        multiplier = self.config.value
        if multiplier >= 1.0:
            return "High Quality (Preserve/Increase Bitrate)"
        elif multiplier >= 0.75:
            return "Medium Quality (Moderate Compression)"
        elif multiplier >= 0.5:
            return "Low Quality (High Compression)"
        else:
            return "Very Low Quality (Maximum Compression)"
    
    @staticmethod
    def get_available_presets() -> Dict[str, Dict]:
        """Get all available quality presets"""
        return {
            'crf': {name: {'value': value, 'description': f"CRF {value}"} 
                   for name, value in EncodingConfigManager.CRF_PRESETS.items()},
            'av1_crf': {name: {'value': value, 'description': f"AV1 CRF {value}"} 
                       for name, value in EncodingConfigManager.AV1_CRF_PRESETS.items()},
            'vbr': {name: {'value': value, 'description': f"{int(value*100)}% of original"} 
                   for name, value in EncodingConfigManager.VBR_PRESETS.items()}
        }


def main():
    """Test encoding configuration functionality"""
    manager = EncodingConfigManager()
    
    print("Encoding Configuration Test")
    print("\nAvailable CRF presets:")
    for name, value in manager.CRF_PRESETS.items():
        print(f"  {name}: CRF {value}")
    
    print("\nAvailable VBR presets:")
    for name, value in manager.VBR_PRESETS.items():
        print(f"  {name}: {int(value*100)}% of original bitrate")
    
    # Test CRF configuration
    print("\n--- Testing CRF Configuration ---")
    manager.set_crf_encoding(23, "medium")
    summary = manager.get_config_summary()
    print(f"Method: {summary['description']}")
    print(f"Codec: {summary['video_codec']}")
    print(f"Quality: {summary['quality_preset']}")
    
    # Test VBR configuration
    print("\n--- Testing VBR Configuration ---")
    manager.set_vbr_encoding(0.75, "fast")
    summary = manager.get_config_summary()
    print(f"Method: {summary['description']}")
    print(f"Quality: {summary['bitrate_preset']}")
    
    # Test bitrate calculation
    original_bitrate = "5000000"  # 5 Mbps
    target_bitrate = manager.calculate_target_bitrate(original_bitrate)
    print(f"Original: {int(original_bitrate)/1000000} Mbps -> Target: {target_bitrate/1000000} Mbps")


if __name__ == "__main__":
    main()