import gradio as gr
import subprocess
import tempfile
import os
from pathlib import Path
import shutil
import json
import logging
import threading
import time

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

def adjust_dimensions(width, height, codec):
    """Adjust dimensions to be compatible with the codec's requirements"""
    if codec in ["libx265", "libvpx-vp9"]:
        # These codecs require dimensions to be multiples of 8
        new_width = (width + 7) // 8 * 8
        new_height = (height + 7) // 8 * 8
        logger.debug(f"Adjusting dimensions from {width}x{height} to {new_width}x{new_height} for codec {codec}")
        return new_width, new_height
    return width, height

def compress_video(
    input_video,
    resolution,
    quality,
    output_format,
    video_codec,
    audio_codec,
    preset,
    tune,
    bitrate,
    framerate,
    audio_bitrate,
    audio_sample_rate,
    start_time,
    end_time,
    rotate,
    remove_metadata,
    two_pass,
    custom_params,
    progress=gr.Progress()
):
    if input_video is None:
        return None, None, "Please upload a video file first."
    
    input_path = None
    output_path = None
    
    try:
        progress(0, desc="Starting compression...")
        logger.debug(f"Starting compression for file: {input_video}")
        
        # Create temporary files
        progress(0.1, desc="Creating temporary files...")
        with tempfile.NamedTemporaryFile(delete=False, suffix=Path(input_video).name) as tmp_input:
            shutil.copy2(input_video, tmp_input.name)
            input_path = tmp_input.name
        
        output_ext = f".{output_format.lower()}"
        output_path = input_path.replace(Path(input_video).suffix, f"_compressed{output_ext}")
        
        logger.debug(f"Input path: {input_path}")
        logger.debug(f"Output path: {output_path}")
        
        # Get video info to check dimensions
        progress(0.2, desc="Analyzing video...")
        info_cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", input_path]
        logger.debug(f"Running ffprobe command: {' '.join(info_cmd)}")
        
        probe_result = subprocess.run(info_cmd, capture_output=True, text=True)
        if probe_result.returncode != 0:
            logger.error(f"ffprobe failed: {probe_result.stderr}")
            return None, None, f"Failed to analyze video: {probe_result.stderr}"
        
        video_info = json.loads(probe_result.stdout)
        logger.debug(f"Video info: {video_info}")
        
        # Extract video stream info
        video_stream = next((s for s in video_info["streams"] if s["codec_type"] == "video"), None)
        needs_resize = False
        scale_filter = ""
        
        if video_stream:
            width = video_stream.get("width", 0)
            height = video_stream.get("height", 0)
            logger.debug(f"Original dimensions: {width}x{height}")
            
            # Resolution scaling
            scale_map = {
                "Original": "",
                "240p": "scale=-1:240",
                "360p": "scale=-1:360",
                "480p": "scale=-1:480",
                "720p": "scale=-1:720",
                "1080p": "scale=-1:1080",
                "1440p": "scale=-1:1440",
                "4K": "scale=-1:2160"
            }
            
            if resolution != "Original":
                scale_filter = scale_map[resolution]
                logger.debug(f"Using resolution scale: {scale_filter}")
            else:
                # Check if we need to adjust dimensions for codec compatibility
                new_width, new_height = adjust_dimensions(width, height, video_codec)
                if new_width != width or new_height != height:
                    scale_filter = f"scale={new_width}:{new_height}"
                    needs_resize = True
                    logger.debug(f"Using dimension adjustment: {scale_filter}")
        else:
            logger.warning("No video stream found")
        
        # Build FFmpeg command
        progress(0.3, desc="Building compression command...")
        cmd = ["ffmpeg"]
        
        # Add input seeking if specified
        if start_time:
            cmd.extend(["-ss", start_time])
        
        cmd.extend(["-i", input_path])
        
        # Add end time if specified
        if end_time:
            cmd.extend(["-to", end_time])
        
        # Video filters
        vf_filters = []
        
        # Add scale filter if needed
        if scale_filter:
            vf_filters.append(scale_filter)
        
        # Rotation
        if rotate != "None":
            rotate_map = {
                "90°": "transpose=1",
                "180°": "transpose=1,transpose=1",
                "270°": "transpose=2",
                "Flip Horizontal": "hflip",
                "Flip Vertical": "vflip"
            }
            vf_filters.append(rotate_map[rotate])
        
        # Add video filters if any
        if vf_filters:
            cmd.extend(["-vf", ",".join(vf_filters)])
            logger.debug(f"Video filters: {','.join(vf_filters)}")
        
        # Frame rate
        if framerate != "Original":
            cmd.extend(["-r", framerate])
            logger.debug(f"Frame rate: {framerate}")
        
        # Video codec settings
        if video_codec == "libx264":
            # Map quality to CRF for H.264
            crf_map = {"Low": 28, "Medium": 23, "High": 18, "Ultra High": 15}
            crf = crf_map[quality]
            cmd.extend(["-c:v", video_codec, "-crf", str(crf)])
            logger.debug(f"H.264 CRF: {crf}")
        elif video_codec == "libx265":
            # Map quality to CRF for H.265 (generally lower values for same quality)
            crf_map = {"Low": 32, "Medium": 28, "High": 24, "Ultra High": 20}
            crf = crf_map[quality]
            cmd.extend(["-c:v", video_codec, "-crf", str(crf)])
            logger.debug(f"H.265 CRF: {crf}")
        elif video_codec == "libvpx-vp9":
            # Map quality to CRF for VP9
            crf_map = {"Low": 31, "Medium": 26, "High": 21, "Ultra High": 15}
            crf = crf_map[quality]
            cmd.extend(["-c:v", video_codec, "-crf", str(crf), "-b:v", "0"])
            logger.debug(f"VP9 CRF: {crf}")
        else:
            cmd.extend(["-c:v", video_codec])
            logger.debug(f"Video codec: {video_codec}")
        
        # Preset and tuning
        if preset != "medium":
            cmd.extend(["-preset", preset])
            logger.debug(f"Preset: {preset}")
        
        if tune != "none":
            cmd.extend(["-tune", tune])
            logger.debug(f"Tune: {tune}")
        
        # Bitrate
        if bitrate:
            cmd.extend(["-b:v", bitrate])
            logger.debug(f"Video bitrate: {bitrate}")
        
        # Two-pass encoding
        if two_pass:
            logger.debug("Using two-pass encoding")
            # First pass
            progress(0.4, desc="First pass encoding...")
            first_pass_cmd = cmd.copy()
            first_pass_cmd.extend(["-pass", "1", "-f", "null", "/dev/null"])
            logger.debug(f"First pass command: {' '.join(first_pass_cmd)}")
            
            result = subprocess.run(first_pass_cmd, capture_output=True, text=True)
            if result.returncode != 0:
                logger.error(f"First pass failed: {result.stderr}")
                return None, None, f"First pass encoding failed: {result.stderr}"
            
            # Second pass
            progress(0.7, desc="Second pass encoding...")
            cmd.extend(["-pass", "2"])
        else:
            progress(0.4, desc="Encoding video...")
        
        # Audio settings
        if audio_codec == "copy":
            cmd.extend(["-c:a", "copy"])
            logger.debug("Copying audio stream")
        else:
            cmd.extend(["-c:a", audio_codec])
            logger.debug(f"Audio codec: {audio_codec}")
            
            if audio_bitrate != "Auto":
                cmd.extend(["-b:a", audio_bitrate])
                logger.debug(f"Audio bitrate: {audio_bitrate}")
            
            if audio_sample_rate != "Original":
                cmd.extend(["-ar", audio_sample_rate])
                logger.debug(f"Audio sample rate: {audio_sample_rate}")
        
        # Remove metadata
        if remove_metadata:
            cmd.extend(["-map_metadata", "-1"])
            logger.debug("Removing metadata")
        
        # Fast start for web playback
        cmd.extend(["-movflags", "+faststart"])
        
        # Custom parameters
        if custom_params:
            try:
                custom_list = json.loads(custom_params)
                cmd.extend(custom_list)
                logger.debug(f"Custom parameters: {custom_list}")
            except Exception as e:
                logger.error(f"Failed to parse custom parameters: {e}")
        
        # Output file
        cmd.extend(["-y", output_path])
        
        logger.debug(f"Final FFmpeg command: {' '.join(cmd)}")
        
        # Run FFmpeg with progress tracking
        progress(0.8, desc="Processing video...")
        
        # Start FFmpeg process
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True
        )
        
        # Track progress
        duration = None
        while True:
            line = process.stderr.readline()
            if not line:
                break
            
            # Parse FFmpeg progress output
            if "Duration" in line:
                try:
                    duration_str = line.split("Duration: ")[1].split(",")[0].strip()
                    h, m, s = duration_str.split(":")
                    duration = float(h) * 3600 + float(m) * 60 + float(s)
                except:
                    pass
            
            if "time=" in line and duration:
                try:
                    time_str = line.split("time=")[1].split(" ")[0].strip()
                    h, m, s = time_str.split(":")
                    current_time = float(h) * 3600 + float(m) * 60 + float(s)
                    progress_value = 0.8 + (current_time / duration) * 0.2
                    progress(progress_value, desc=f"Processing: {int(progress_value * 100)}%")
                except:
                    pass
        
        # Wait for process to complete
        return_code = process.wait()
        if return_code != 0:
            stderr = process.stderr.read()
            raise subprocess.CalledProcessError(return_code, cmd, stderr)
        
        logger.debug("FFmpeg command completed successfully")
        
        # Clean up pass files if two-pass was used
        if two_pass:
            for f in ["ffmpeg2pass-0.log", "ffmpeg2pass-0.log.mbtree"]:
                if os.path.exists(f):
                    os.unlink(f)
        
        progress(0.95, desc="Finalizing...")
        # Calculate file sizes
        original_size = os.path.getsize(input_path) / (1024 * 1024)
        compressed_size = os.path.getsize(output_path) / (1024 * 1024)
        compression_ratio = (1 - compressed_size / original_size) * 100
        
        # Get video info
        info_cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", output_path]
        probe_result = subprocess.run(info_cmd, capture_output=True, text=True)
        video_info = json.loads(probe_result.stdout)
        
        # Extract relevant info
        video_stream = next((s for s in video_info["streams"] if s["codec_type"] == "video"), None)
        audio_stream = next((s for s in video_info["streams"] if s["codec_type"] == "audio"), None)
        
        # Create filename for download
        original_name = Path(input_video).stem
        download_name = f"{original_name}_compressed.{output_format.lower()}"
        
        info_text = f"""
        **✅ Compression Complete!**
        
        **File Sizes:**
        - **Original:** {original_size:.2f} MB
        - **Compressed:** {compressed_size:.2f} MB
        - **Compression Ratio:** {compression_ratio:.1f}%
        
        **Video Properties:**
        - **Resolution:** {video_stream.get('width', 'N/A')}x{video_stream.get('height', 'N/A')}
        - **Codec:** {video_stream.get('codec_name', 'N/A')}
        - **Frame Rate:** {video_stream.get('r_frame_rate', 'N/A')}
        - **Bitrate:** {int(video_stream.get('bit_rate', 0))/1000:.0f} kbps
        
        **Audio Properties:**
        - **Codec:** {audio_stream.get('codec_name', 'N/A') if audio_stream else 'N/A'}
        - **Sample Rate:** {audio_stream.get('sample_rate', 'N/A') if audio_stream else 'N/A'} Hz
        - **Channels:** {audio_stream.get('channels', 'N/A') if audio_stream else 'N/A'}
        
        **Settings Used:**
        - Resolution: {resolution}
        - Quality: {quality}
        - Output Format: {output_format}
        - Video Codec: {video_codec}
        - Audio Codec: {audio_codec}
        - Preset: {preset}
        - Tune: {tune}
        """
        
        progress(1.0, desc="Done!")
        return output_path, download_name, info_text
        
    except subprocess.CalledProcessError as e:
        error_message = e.stderr
        logger.error(f"FFmpeg error: {error_message}")
        
        # Check for specific error conditions
        if "Picture width must be an integer multiple" in error_message:
            error_message = "The video dimensions are not compatible with the selected codec. Try using H.264 instead."
        elif "Cannot open libx265 encoder" in error_message:
            error_message = "Error with H.265 encoding. Try using H.264 instead."
        elif "Unsupported codec" in error_message:
            error_message = "The selected codec is not supported. Try a different codec."
        
        return None, None, f"❌ Compression failed: {error_message}"
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        return None, None, f"❌ An unexpected error occurred: {str(e)}"
    finally:
        # Clean up input file
        if input_path and os.path.exists(input_path):
            os.unlink(input_path)
            logger.debug("Cleaned up input file")

# Create the Gradio interface
with gr.Blocks(title="Advanced Video Compressor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 Advanced Video Compressor")
    gr.Markdown("Upload a video file and configure advanced compression settings.")
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input section
            input_video = gr.File(
                label="Upload Video",
                file_types=["video"],
                file_count="single"
            )
            
            # Basic settings
            with gr.Group():
                gr.Markdown("### 🎯 Basic Settings")
                
                with gr.Row():
                    resolution = gr.Dropdown(
                        choices=["Original", "240p", "360p", "480p", "720p", "1080p", "1440p", "4K"],
                        value="Original",
                        label="Resolution"
                    )
                    
                    quality = gr.Dropdown(
                        choices=["Low", "Medium", "High", "Ultra High"],
                        value="Medium",
                        label="Quality"
                    )
                
                with gr.Row():
                    output_format = gr.Dropdown(
                        choices=["MP4", "MKV", "WEBM", "AVI", "MOV"],
                        value="MP4",
                        label="Output Format"
                    )
                    
                    rotate = gr.Dropdown(
                        choices=["None", "90°", "180°", "270°", "Flip Horizontal", "Flip Vertical"],
                        value="None",
                        label="Rotate/Flip"
                    )
            
            # Video settings
            with gr.Group():
                gr.Markdown("### 🎥 Video Settings")
                
                with gr.Row():
                    video_codec = gr.Dropdown(
                        choices=["libx264", "libx265", "libvpx-vp9", "mpeg4"],
                        value="libx264",
                        label="Video Codec"
                    )
                    
                    preset = gr.Dropdown(
                        choices=["ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow"],
                        value="medium",
                        label="Encoding Preset"
                    )
                
                with gr.Row():
                    tune = gr.Dropdown(
                        choices=["none", "film", "animation", "grain", "stillimage", "psnr", "ssim", "fastdecode", "zerolatency"],
                        value="none",
                        label="Tuning"
                    )
                    
                    framerate = gr.Dropdown(
                        choices=["Original", "15", "24", "25", "30", "60"],
                        value="Original",
                        label="Frame Rate (fps)"
                    )
                
                bitrate = gr.Textbox(
                    label="Video Bitrate (e.g., 1000K, 1M)",
                    placeholder="Leave empty for automatic",
                    value=""
                )
                
                two_pass = gr.Checkbox(
                    label="Two-Pass Encoding (better quality)",
                    value=False
                )
            
            # Audio settings
            with gr.Group():
                gr.Markdown("### 🎵 Audio Settings")
                
                with gr.Row():
                    audio_codec = gr.Dropdown(
                        choices=["aac", "libmp3lame", "libvorbis", "opus", "copy"],
                        value="aac",
                        label="Audio Codec"
                    )
                    
                    audio_bitrate = gr.Dropdown(
                        choices=["Auto", "32k", "64k", "96k", "128k", "192k", "256k", "320k"],
                        value="Auto",
                        label="Audio Bitrate"
                    )
                
                audio_sample_rate = gr.Dropdown(
                    choices=["Original", "22050", "44100", "48000"],
                    value="Original",
                    label="Sample Rate (Hz)"
                )
            
            # Time settings
            with gr.Group():
                gr.Markdown("### ⏱️ Time Settings")
                
                with gr.Row():
                    start_time = gr.Textbox(
                        label="Start Time (HH:MM:SS)",
                        placeholder="00:00:00",
                        value=""
                    )
                    
                    end_time = gr.Textbox(
                        label="End Time (HH:MM:SS)",
                        placeholder="00:00:00",
                        value=""
                    )
            
            # Advanced settings
            with gr.Group():
                gr.Markdown("### ⚙️ Advanced Settings")
                
                remove_metadata = gr.Checkbox(
                    label="Remove Metadata",
                    value=False
                )
                
                custom_params = gr.Textbox(
                    label="Custom FFmpeg Parameters (JSON array)",
                    placeholder='["-vf", "scale=1280:720"]',
                    value=""
                )
            
            # Compress button
            compress_btn = gr.Button("🚀 Compress Video", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            # Progress bar
            progress_html = gr.HTML(
                value='<div class="progress-container"><div class="progress-bar" style="width: 0%"></div></div><div class="progress-text">Ready to compress</div>'
            )
            
            # Output section
            with gr.Group():
                gr.Markdown("### 📥 Download")
                
                output_video = gr.File(
                    label="Compressed Video",
                    visible=False
                )
                
                download_btn = gr.DownloadButton(
                    label="⬇️ Download Compressed Video",
                    visible=False,
                    size="lg"
                )
            
            # Info section
            with gr.Group():
                gr.Markdown("### 📊 Compression Info")
                
                output_info = gr.Markdown(
                    value="Upload a video and configure settings, then click 'Compress Video' to begin."
                )
    
    # Custom CSS for progress bar
    gr.HTML("""
    <style>
    .progress-container {
        width: 100%;
        background-color: #f0f0f0;
        border-radius: 4px;
        margin-bottom: 10px;
    }
    .progress-bar {
        height: 20px;
        background-color: #007bff;
        border-radius: 4px;
        transition: width 0.3s ease;
    }
    .progress-text {
        text-align: center;
        font-size: 14px;
        color: #666;
    }
    </style>
    """)
    
    # Set up the compression function
    def compress_with_progress(*args):
        # Create a progress callback that updates the HTML
        def update_progress(value, desc=""):
            progress_html.update(f'''
                <div class="progress-container">
                    <div class="progress-bar" style="width: {value * 100}%"></div>
                </div>
                <div class="progress-text">{desc} ({int(value * 100)}%)</div>
            ''')
        
        # Call the actual compression function with our progress callback
        return compress_video(*args[:-1], progress=update_progress)
    
    compress_btn.click(
        fn=compress_with_progress,
        inputs=[
            input_video, resolution, quality, output_format, video_codec, audio_codec,
            preset, tune, bitrate, framerate, audio_bitrate, audio_sample_rate,
            start_time, end_time, rotate, remove_metadata, two_pass, custom_params, progress_html
        ],
        outputs=[output_video, download_btn, output_info]
    )
    
    # Auto-hide output when input changes
    input_video.change(
        fn=lambda x: (gr.update(visible=False), gr.update(visible=False), gr.update(value='<div class="progress-container"><div class="progress-bar" style="width: 0%"></div></div><div class="progress-text">Ready to compress</div>'), "Upload a video and configure settings, then click 'Compress Video' to begin."),
        inputs=[input_video],
        outputs=[output_video, download_btn, progress_html, output_info]
    )

if __name__ == "__main__":
    demo.launch()