import os
import subprocess
import shutil
from pathlib import Path
import sys
import time
import json

# Add paths for ffmpeg and ffprobe
FFMPEG_PATH = '/opt/homebrew/bin/ffmpeg'
FFPROBE_PATH = '/opt/homebrew/bin/ffprobe'

# Target maximum file size (in bytes)
TARGET_SIZE_MB = 50
TARGET_SIZE_BYTES = TARGET_SIZE_MB * 1024 * 1024

# 快速模式（单遍编码，速度快但质量略低）
FAST_MODE = True  # 设置为True启用快速模式

# Compression presets
COMPRESSION_LEVELS = {
    "moderate": {
        "resolution_scale": 1.0,  # No scaling
        "crf": 28,  # Higher CRF = more compression, lower quality
        "preset": "medium"
    },
    "high": {
        "resolution_scale": 0.75,  # Scale to 75% of original resolution
        "crf": 32,  # 提高CRF值以获得更高压缩
        "preset": "slow"  # Better compression
    },
    "extreme": {
        "resolution_scale": 0.5,  # Scale to 50% of original resolution
        "crf": 35,  # 更高的CRF值
        "preset": "slow"
    },
    "ultra": {  # 新增超高压缩级别
        "resolution_scale": 0.5,
        "crf": 38,
        "preset": "veryslow"
    }
}


def get_video_info(video_path):
    """Get video information using ffprobe."""
    cmd = [
        FFPROBE_PATH,  # Use the full path to ffprobe
        '-v', 'quiet',
        '-print_format', 'json',
        '-show_format',
        '-show_streams',
        video_path
    ]
    try:
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode != 0:
            print(f"Error getting video info for {video_path}: {result.stderr}")
            return None
        
        # Parse the JSON output to get video information
        video_info = json.loads(result.stdout)
        return video_info
    except Exception as e:
        print(f"Exception when getting video info: {e}")
        return None


def compress_video(input_path, output_path, target_size_bytes=TARGET_SIZE_BYTES, compression_level="moderate"):
    """Compress video to target size while maintaining reasonable quality."""
    # Get file size
    file_size = os.path.getsize(input_path)

    # If file is already smaller than target, just copy it
    if file_size <= target_size_bytes:
        print(f"{input_path} is already smaller than {TARGET_SIZE_MB}MB, copying...")
        shutil.copy2(input_path, output_path)
        return True

    # Get video info to calculate dimensions
    video_info = get_video_info(input_path)
    if not video_info:
        print(f"Could not get video information for {input_path}")
        return False
    
    # Extract video dimensions and duration
    video_stream = None
    for stream in video_info.get('streams', []):
        if stream.get('codec_type') == 'video':
            video_stream = stream
            break
    
    if not video_stream:
        print(f"No video stream found in {input_path}")
        return False
    
    # Get original width and height
    width = int(video_stream.get('width', 1280))
    height = int(video_stream.get('height', 720))
    
    # Get duration from format section (more reliable)
    try:
        duration = float(video_info.get('format', {}).get('duration', 0))
        if duration == 0:
            # Fallback to stream duration
            duration = float(video_stream.get('duration', 0))
    except (ValueError, TypeError):
        # Try to get duration using ffprobe command as fallback
        cmd_duration = [
            FFPROBE_PATH,  # Use FFPROBE_PATH here too
            '-v', 'error',
            '-show_entries', 'format=duration',
            '-of', 'default=noprint_wrappers=1:nokey=1',
            input_path
        ]
        duration = float(subprocess.check_output(cmd_duration).decode('utf-8').strip())

    # Get compression settings
    comp_settings = COMPRESSION_LEVELS.get(compression_level, COMPRESSION_LEVELS["moderate"])
    
    # Use fixed dimensions 645x360
    new_width = 645
    new_height = 360
    
    # Ensure dimensions are even (required by some codecs)
    new_width = new_width - (new_width % 2)  # 645 will become 644
    new_height = new_height - (new_height % 2)  # 360 remains 360
    
    # Calculate target bitrate based on target size and duration
    # Use a lower percentage (85%) to account for container overhead
    target_bitrate = int((target_size_bytes * 8 * 0.85) / duration)
    
    # Set minimum bitrate
    min_bitrate = 50000  # 50kbps
    if target_bitrate < min_bitrate:
        target_bitrate = min_bitrate
        print(f"Warning: Target bitrate too low, using minimum of {min_bitrate/1000}kbps")

    print(f"Compressing {input_path} (current size: {file_size / 1024 / 1024:.2f}MB)")
    print(f"Original resolution: {width}x{height}, New resolution: {new_width}x{new_height}")
    print(f"Compression level: {compression_level}, CRF: {comp_settings['crf']}")
    print(f"Target bitrate: {target_bitrate / 1000:.2f}kbps")

    # Ensure target directory exists
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # 快速模式：使用单遍编码
    if FAST_MODE:
        print("使用快速模式（单遍编码）...")
        cmd = [
            FFMPEG_PATH,
            '-y',  # Overwrite output files
            '-i', input_path,
            '-c:v', 'libx264',
            '-preset', 'fast',  # 使用fast预设以加快速度
            '-crf', str(comp_settings['crf']),
            '-vf', f'scale={new_width}:{new_height}',
            '-b:v', f'{target_bitrate}',
            '-maxrate', f'{int(target_bitrate * 1.5)}',
            '-bufsize', f'{int(target_bitrate * 2)}',
            '-c:a', 'aac',
            '-b:a', '48k',
            '-ac', '2',
            output_path
        ]
        
        print("运行压缩...")
        process = subprocess.Popen(
            cmd, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE, 
            universal_newlines=True,
            bufsize=1
        )
        
        stderr_lines = []
        _display_ffmpeg_progress(process, duration, stderr_lines)
        result = process.wait()
        
        if result != 0:
            error_details = "".join(stderr_lines)
            print(f"\nError compressing video: ffmpeg exited with code {result}")
            print(f"Error details: {error_details if error_details else 'No error details available'}")
            return False
    else:
        # 原有的两遍编码逻辑
        # Use two-pass encoding for better quality
        temp_folder = os.path.dirname(output_path)
        null_path = "NUL" if os.name == "nt" else "/dev/null"
        
        # Create a unique stats file path that doesn't contain special characters
        stats_file = os.path.join(temp_folder, f"ffmpeg_stats_{int(time.time())}")

        # CRF-based encoding first pass
        cmd_pass1 = [
            FFMPEG_PATH,  # Use FFMPEG_PATH instead of 'ffmpeg'
            '-y',  # Overwrite output files
            '-i', input_path,
            '-c:v', 'libx264',
            '-preset', comp_settings['preset'],
            '-crf', str(comp_settings['crf']),
            '-vf', f'scale={new_width}:{new_height}',
            '-pass', '1',
            '-f', 'mp4',
            '-passlogfile', stats_file,
            '-an',  # No audio in first pass
            null_path
        ]

        # Add bitrate constraint for second pass
        cmd_pass2 = [
            FFMPEG_PATH,  # Use FFMPEG_PATH instead of 'ffmpeg'
            '-y',  # Overwrite output files
            '-i', input_path,
            '-c:v', 'libx264',
            '-preset', comp_settings['preset'],
            '-crf', str(comp_settings['crf']),
            '-vf', f'scale={new_width}:{new_height}',
            '-b:v', f'{target_bitrate}',
            '-maxrate', f'{int(target_bitrate * 1.5)}',
            '-bufsize', f'{int(target_bitrate * 2)}', 
            '-pass', '2',
            '-passlogfile', stats_file,
            '-c:a', 'aac',
            '-b:a', '48k',  # 降低音频比特率到48k
            '-ac', '2',     # Stereo audio
            output_path
        ]

        print("Running first pass...")
        process = subprocess.Popen(
            cmd_pass1, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE, 
            universal_newlines=True,
            bufsize=1
        )
        
        # Show progress for first pass
        _display_ffmpeg_progress(process, duration)
        first_pass_result = process.wait()
        
        if first_pass_result != 0:
            stderr_output = process.stderr.read()
            print(f"\nError in first pass: ffmpeg exited with code {first_pass_result}")
            print(f"Error details: {stderr_output}")
            
            # Clean up any existing pass log files before returning
            for file in Path(temp_folder).glob(f"{os.path.basename(stats_file)}*"):
                try:
                    file.unlink()
                except:
                    pass
                
            return False

        print("\nRunning second pass...")
        process = subprocess.Popen(
            cmd_pass2, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE, 
            universal_newlines=True,
            bufsize=1
        )
        
        # Capture error output
        stderr_lines = []
        
        # Show progress for second pass
        try:
            _display_ffmpeg_progress(process, duration, stderr_lines)
            result = process.wait()
        except Exception as e:
            print(f"\nError during progress display: {e}")
            # Ensure process is terminated
            process.kill()
            result = -1
        
        # Clean up pass log files
        for file in Path(temp_folder).glob(f"{os.path.basename(stats_file)}*"):
            try:
                file.unlink()
            except:
                pass

        if result != 0:
            error_details = "".join(stderr_lines)
            print(f"\nError compressing video: ffmpeg exited with code {result}")
            print(f"Error details: {error_details if error_details else 'No error details available'}")
            
            # Try to get any remaining error output
            if process.stderr:
                remaining_stderr = process.stderr.read()
                if remaining_stderr:
                    print(f"Additional error output: {remaining_stderr}")
            return False

    # Verify the compressed file exists and has a reasonable size
    if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
        new_size = os.path.getsize(output_path) / 1024 / 1024
        compression_ratio = file_size / os.path.getsize(output_path)
        print(f"Compression complete. New size: {new_size:.2f}MB (Reduced by {(1-new_size/file_size*100):.1f}%, ratio 1:{compression_ratio:.1f})")
        return True
    else:
        print("Compression failed: Output file is missing or empty")
        return False


def _display_ffmpeg_progress(process, total_duration, stderr_buffer=None):
    """Display progress for ffmpeg process."""
    progress_length = 40  # Length of the progress bar
    
    # Pattern to look for in ffmpeg stderr output
    time_pattern = "time="
    
    for line in process.stderr:
        # Save for error reporting if needed
        if stderr_buffer is not None:
            stderr_buffer.append(line)
            
        if time_pattern in line:
            # Extract the current time
            time_index = line.find(time_pattern)
            time_str = line[time_index + len(time_pattern):].split(' ')[0]
            
            # Skip if time is N/A
            if time_str == 'N/A':
                continue
                
            # Convert the time string (HH:MM:SS.MS) to seconds
            if ':' in time_str:
                h, m, s = time_str.split(':')
                current_time = float(h) * 3600 + float(m) * 60 + float(s)
            else:
                current_time = float(time_str)
            
            # Calculate progress percentage
            progress = min(1, current_time / total_duration)
            filled_length = int(progress_length * progress)
            
            # Create the progress bar
            bar = '█' * filled_length + '-' * (progress_length - filled_length)
            percentage = progress * 100
            
            # Write progress to stdout
            sys.stdout.write(f"\r[{bar}] {percentage:.1f}% ({current_time:.1f}s/{total_duration:.1f}s)")
            sys.stdout.flush()


def process_folder(folder_path):
    """Process all MP4 files in the given folder and its subfolders."""
    # Create a temporary directory for compressed files
    temp_dir = os.path.join(folder_path, "temp_compressed")
    os.makedirs(temp_dir, exist_ok=True)

    mp4_files = []

    # Find all MP4 files
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith('.mp4'):
                mp4_files.append(os.path.join(root, file))

    if not mp4_files:
        print(f"No MP4 files found in {folder_path}")
        return

    # Check for FFmpeg installation
    if not get_video_info(mp4_files[0]):
        print("Error: FFmpeg is not installed or not working properly.")
        print("Please install FFmpeg and make sure it's available in your PATH.")
        return
        
    # Count total files and calculate total size
    total_files = len(mp4_files)
    total_original_size = sum(os.path.getsize(f) for f in mp4_files) / (1024 * 1024)
    print(f"Found {total_files} MP4 files, total size: {total_original_size:.2f}MB")

    # Process each file
    successful_compressions = 0
    total_new_size = 0
    
    for index, file_path in enumerate(mp4_files, 1):
        print(f"\nProcessing file {index}/{total_files}: {file_path}")
        file_size_mb = os.path.getsize(file_path) / 1024 / 1024
        print(f"Current size: {file_size_mb:.2f}MB")

        if file_size_mb <= TARGET_SIZE_MB:
            print(f"File is already smaller than {TARGET_SIZE_MB}MB, skipping...")
            total_new_size += file_size_mb
            continue

        # Create output path in temp directory
        rel_path = os.path.relpath(file_path, folder_path)
        temp_output_path = os.path.join(temp_dir, rel_path)
        os.makedirs(os.path.dirname(temp_output_path), exist_ok=True)

        # Try progressive compression levels until target is reached
        for level in ["moderate", "high", "extreme", "ultra"]:
            print(f"Trying compression level: {level}")
            
            # Compress the video
            success = compress_video(file_path, temp_output_path, TARGET_SIZE_BYTES, level)
            
            if not success:
                print(f"Compression failed at level {level}")
                continue
                
            new_size_mb = os.path.getsize(temp_output_path) / 1024 / 1024
            
            # If we reached target size or we're on extreme level, we're done
            if new_size_mb <= TARGET_SIZE_MB or level == "ultra":
                break
            
            print(f"Size after {level} compression: {new_size_mb:.2f}MB, still above target {TARGET_SIZE_MB}MB")
            print("Trying higher compression level...")

        if success:
            # Check if the new file is valid
            new_size_mb = os.path.getsize(temp_output_path) / 1024 / 1024
            total_new_size += new_size_mb
            
            # Create temporary backup in case of failure
            backup_path = file_path + ".bak"
            print(f"Creating temporary backup at {backup_path}")
            shutil.copy2(file_path, backup_path)

            # Replace original with compressed version
            print(f"Replacing original with compressed version ({new_size_mb:.2f}MB)")
            shutil.copy2(temp_output_path, file_path)
            
            # Delete the backup file after successful replacement
            print(f"Deleting backup file: {backup_path}")
            os.remove(backup_path)
            
            successful_compressions += 1
        else:
            print("All compression attempts failed, keeping original")
            total_new_size += file_size_mb

    # Clean up temp directory
    shutil.rmtree(temp_dir)
    
    # Print summary
    print("\nCompression Summary:")
    print(f"Processed {total_files} files, successfully compressed {successful_compressions} files")
    print(f"Original total size: {total_original_size:.2f}MB")
    print(f"New total size: {total_new_size:.2f}MB")
    print(f"Space saved: {total_original_size - total_new_size:.2f}MB ({(1-total_new_size/total_original_size)*100:.1f}%)")
    print("\nProcessing complete!")


if __name__ == "__main__":
    test_folder = os.path.join(os.path.dirname(os.path.abspath(__file__)), "2023年GPT夏季班")

    if not os.path.exists(test_folder):
        print(f"Error: Test folder '{test_folder}' not found.")
    else:
        print(f"Starting compression of MP4 files in {test_folder}")
        print(f"Target size: {TARGET_SIZE_MB}MB")
        process_folder(test_folder)