"""
Video assembler module for concatenating processed segments.
Handles merging processed segments and adding audio from original video.
"""

import os
import subprocess
import logging
from typing import List, Optional
from utils.file_utils import get_segments_directories, file_exists

logger = logging.getLogger(__name__)

class VideoAssembler:
    """Handles final video assembly from processed segments."""
    
    def __init__(self, preserve_audio: bool = True, use_nvenc: bool = False, 
                 output_mode: str = "green_screen"):
        """
        Initialize video assembler.
        
        Args:
            preserve_audio: Whether to preserve audio from original video
            use_nvenc: Whether to use hardware encoding for final output
            output_mode: Output mode - "green_screen" or "alpha_channel"
        """
        self.preserve_audio = preserve_audio
        self.use_nvenc = use_nvenc
        self.output_mode = output_mode
    
    def create_concat_file(self, segments_dir: str, output_filename: str = "concat_list.txt") -> Optional[str]:
        """
        Create a concatenation file for FFmpeg.
        
        Args:
            segments_dir: Directory containing processed segments
            output_filename: Name for the concat file
            
        Returns:
            Path to concat file or None if no valid segments found
        """
        concat_path = os.path.join(segments_dir, output_filename)
        valid_segments = 0
        
        try:
            segments = get_segments_directories(segments_dir)
            
            with open(concat_path, 'w') as f:
                for i, segment in enumerate(segments):
                    segment_dir = os.path.join(segments_dir, segment)
                    if self.output_mode == "alpha_channel":
                        output_video = os.path.join(segment_dir, f"output_{i}.mov")
                    else:
                        output_video = os.path.join(segment_dir, f"output_{i}.mp4")
                    
                    if file_exists(output_video):
                        # Use relative path for FFmpeg
                        relative_path = os.path.relpath(output_video, segments_dir)
                        f.write(f"file '{relative_path}'\n")
                        valid_segments += 1
                    else:
                        logger.warning(f"Output video not found for segment {i}: {output_video}")
            
            if valid_segments == 0:
                logger.error("No valid output segments found for concatenation")
                os.remove(concat_path)
                return None
            
            logger.info(f"Created concatenation file with {valid_segments} segments: {concat_path}")
            return concat_path
            
        except Exception as e:
            logger.error(f"Error creating concatenation file: {e}")
            return None
    
    def concatenate_segments(self, segments_dir: str, output_path: str, 
                           bitrate: str = "50M") -> bool:
        """
        Concatenate video segments using FFmpeg.
        
        Args:
            segments_dir: Directory containing processed segments
            output_path: Path for final concatenated video
            bitrate: Output video bitrate
            
        Returns:
            True if successful
        """
        # Create concatenation file
        concat_file = self.create_concat_file(segments_dir)
        if not concat_file:
            return False
        
        try:
            # Build FFmpeg command
            if self.output_mode == "alpha_channel":
                # For alpha channel, we need to maintain the ProRes codec
                cmd = [
                    'ffmpeg',
                    '-y',  # Overwrite output
                    '-f', 'concat',
                    '-safe', '0',
                    '-i', concat_file,
                    '-c:v', 'copy',  # Copy video codec to preserve alpha
                    '-an',  # No audio for now
                    output_path
                ]
            else:
                cmd = [
                    'ffmpeg',
                    '-y',  # Overwrite output
                    '-f', 'concat',
                    '-safe', '0',
                    '-i', concat_file,
                    '-c:v', 'copy',  # Copy video codec (no re-encoding)
                    '-an',  # No audio for now
                    output_path
                ]
                
                # Use hardware encoding if requested
                if self.use_nvenc:
                    import sys
                    if sys.platform == 'darwin':
                        encoder = 'hevc_videotoolbox'
                    else:
                        encoder = 'hevc_nvenc'
                    
                    # Re-encode with hardware acceleration
                    cmd = [
                        'ffmpeg',
                        '-y',
                        '-f', 'concat',
                        '-safe', '0',
                        '-i', concat_file,
                        '-c:v', encoder,
                        '-preset', 'slow',
                        '-b:v', bitrate,
                        '-pix_fmt', 'yuv420p',
                        '-an',
                        output_path
                    ]
            
            logger.info(f"Running concatenation command: {' '.join(cmd)}")
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                logger.error(f"FFmpeg concatenation failed: {result.stderr}")
                return False
            
            logger.info(f"Successfully concatenated segments to: {output_path}")
            
            # Clean up concat file
            try:
                os.remove(concat_file)
            except:
                pass
            
            return True
            
        except Exception as e:
            logger.error(f"Error during concatenation: {e}")
            return False
    
    def copy_audio_from_original(self, original_video: str, processed_video: str, 
                               final_output: str) -> bool:
        """
        Copy audio track from original video to processed video.
        
        Args:
            original_video: Path to original video with audio
            processed_video: Path to processed video without audio
            final_output: Path for final output with audio
            
        Returns:
            True if successful
        """
        if not self.preserve_audio:
            logger.info("Audio preservation disabled, skipping audio copy")
            return True
        
        try:
            # Check if original video has audio
            probe_cmd = [
                'ffprobe',
                '-v', 'error',
                '-select_streams', 'a:0',
                '-show_entries', 'stream=codec_type',
                '-of', 'csv=p=0',
                original_video
            ]
            
            result = subprocess.run(probe_cmd, capture_output=True, text=True)
            
            if result.returncode != 0 or result.stdout.strip() != 'audio':
                logger.warning("Original video has no audio track")
                # Just copy the processed video
                import shutil
                shutil.copy2(processed_video, final_output)
                return True
            
            # Copy audio from original to processed video
            cmd = [
                'ffmpeg',
                '-y',
                '-i', processed_video,  # Video input
                '-i', original_video,   # Audio input
                '-c:v', 'copy',         # Copy video stream
                '-c:a', 'copy',         # Copy audio stream
                '-map', '0:v:0',        # Map video from first input
                '-map', '1:a:0',        # Map audio from second input
                '-shortest',            # Match duration to shortest stream
                final_output
            ]
            
            logger.info("Copying audio from original video...")
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                logger.error(f"FFmpeg audio copy failed: {result.stderr}")
                return False
            
            logger.info(f"Successfully added audio to final video: {final_output}")
            return True
            
        except Exception as e:
            logger.error(f"Error copying audio: {e}")
            return False
    
    def assemble_final_video(self, segments_dir: str, original_video: str, 
                           output_path: str, bitrate: str = "50M") -> bool:
        """
        Complete pipeline to assemble final video with audio.
        
        Args:
            segments_dir: Directory containing processed segments
            original_video: Path to original video (for audio)
            output_path: Path for final output video
            bitrate: Output video bitrate
            
        Returns:
            True if successful
        """
        logger.info("Starting final video assembly...")
        
        # Step 1: Concatenate segments
        temp_concat_path = os.path.join(os.path.dirname(output_path), "temp_concat.mp4")
        
        if not self.concatenate_segments(segments_dir, temp_concat_path, bitrate):
            logger.error("Failed to concatenate segments")
            return False
        
        # Step 2: Add audio from original
        if self.preserve_audio and file_exists(original_video):
            success = self.copy_audio_from_original(original_video, temp_concat_path, output_path)
            
            # Clean up temp file
            try:
                os.remove(temp_concat_path)
            except:
                pass
                
            return success
        else:
            # No audio to add, just rename temp file
            import shutil
            try:
                shutil.move(temp_concat_path, output_path)
                logger.info(f"Final video saved to: {output_path}")
                return True
            except Exception as e:
                logger.error(f"Error moving final video: {e}")
                return False
    
    def verify_segment_completeness(self, segments_dir: str) -> tuple[bool, List[int]]:
        """
        Verify all segments have been processed.
        
        Args:
            segments_dir: Directory containing segments
            
        Returns:
            Tuple of (all_complete, missing_segments)
        """
        segments = get_segments_directories(segments_dir)
        missing_segments = []
        
        for i, segment in enumerate(segments):
            segment_dir = os.path.join(segments_dir, segment)
            if self.output_mode == "alpha_channel":
                output_video = os.path.join(segment_dir, f"output_{i}.mov")
            else:
                output_video = os.path.join(segment_dir, f"output_{i}.mp4")
            
            if not file_exists(output_video):
                missing_segments.append(i)
        
        all_complete = len(missing_segments) == 0
        
        if all_complete:
            logger.info(f"All {len(segments)} segments have been processed")
        else:
            logger.warning(f"Missing output for segments: {missing_segments}")
        
        return all_complete, missing_segments