import subprocess
import json
import os
import glob

def get_video_duration(video_path):
    """Uses ffprobe to extract the exact duration of the baseline video."""
    print(f"Analyzing baseline duration: {video_path}...")
    cmd = [
        'ffprobe', '-v', 'error', '-show_entries', 'format=duration', 
        '-of', 'default=noprint_wrappers=1:nokey=1', video_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    try:
        duration = float(result.stdout.strip())
        print(f"Baseline duration found: {duration} seconds.")
        return duration
    except ValueError:
        print("Error: Could not determine the duration of the baseline video.")
        return None

def get_next_keyframe(input_video, cut_timestamp):
    """Finds the next keyframe immediately following the cut timestamp."""
    cmd = [
        'ffprobe', '-v', 'quiet', '-select_streams', 'v',
        '-skip_frame', 'nokey', '-show_frames', 
        '-show_entries', 'frame=pkt_pts_time,pkt_dts_time,best_effort_timestamp_time',
        '-of', 'json', '-read_intervals', f'{cut_timestamp}%+10', input_video
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    frames = json.loads(result.stdout).get('frames', [])
    
    for frame in frames:
        # Safely try multiple timestamp keys since different codecs/containers vary
        time_str = (frame.get('best_effort_timestamp_time') or 
                    frame.get('pkt_pts_time') or 
                    frame.get('pkt_dts_time'))
        
        if time_str is not None:
            keyframe_time = float(time_str)
            if keyframe_time > cut_timestamp:
                return keyframe_time
    return None

def smart_cut(input_video, output_video, cut_timestamp):
    """Executes the Smart Cut process: Re-encode a tiny segment, copy the rest, and stitch."""
    print(f"Processing: {input_video}...")
    next_keyframe = get_next_keyframe(input_video, cut_timestamp)
    
    if not next_keyframe:
        print(f"Skipping {input_video}: Could not find a keyframe after {cut_timestamp}s.")
        return

    part1 = f"temp_part1_{input_video}"
    part2 = f"temp_part2_{input_video}"
    concat_list = f"concat_{input_video}.txt"

    try:
        # 1. Re-encode the tiny segment (from exact cut to next keyframe)
        subprocess.run([
            'ffmpeg', '-y', '-v', 'error', '-i', input_video, 
            '-ss', str(cut_timestamp), '-to', str(next_keyframe), 
            '-c:v', 'libx264', '-crf', '18', '-c:a', 'copy', part1
        ], check=True)
        
        # 2. Copy the rest of the video (from next keyframe to the end)
        subprocess.run([
            'ffmpeg', '-y', '-v', 'error', '-ss', str(next_keyframe), '-i', input_video, 
            '-c:v', 'copy', '-c:a', 'copy', part2
        ], check=True)
        
        # 3. Concatenate them using a text list
        with open(concat_list, 'w', encoding='utf-8') as f:
            f.write(f"file '{part1}'\nfile '{part2}'\n")
            
        subprocess.run([
            'ffmpeg', '-y', '-v', 'error', '-f', 'concat', '-safe', '0', 
            '-i', concat_list, '-c', 'copy', output_video
        ], check=True)
        print(f"Success! Saved to {output_video}")

    except subprocess.CalledProcessError as e:
        print(f"Error processing {input_video}. FFmpeg failed.")
    finally:
        # 4. Clean up temporary files safely
        for temp_file in [part1, part2, concat_list]:
            if os.path.exists(temp_file):
                os.remove(temp_file)

def main():
    baseline_file = "baseline.mp4"
    output_dir = "processed_videos"

    # Check if baseline exists
    if not os.path.exists(baseline_file):
        print(f"Error: '{baseline_file}' not found in the current directory.")
        return

    # Get exact cut timestamp from the baseline
    cut_timestamp = get_video_duration(baseline_file)
    if not cut_timestamp:
        return

    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    # Find all mp4 files in the folder
    video_files = glob.glob("*.mp4")
    
    # Remove the baseline file from the processing list
    if baseline_file in video_files:
        video_files.remove(baseline_file)

    if not video_files:
        print("No other .mp4 files found to process.")
        return

    print(f"Found {len(video_files)} videos to process.\n" + "-"*30)

    # Process each video
    for video in video_files:
        output_path = os.path.join(output_dir, video)
        smart_cut(video, output_path, cut_timestamp)

    print("-" * 30 + "\nBatch processing complete!")

if __name__ == "__main__":
    main()