import subprocess
import json
import os
import glob
import sys

def get_video_duration(video_path):
    """Uses ffprobe to extract the exact duration of a video."""
    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:
        return float(result.stdout.strip())
    except ValueError:
        return None

def get_previous_keyframe(input_video, cut_timestamp):
    """Scans backward to find the closest keyframe immediately PRECEDING the cut timestamp."""
    # Look up to 5 minutes backward for a keyframe to be safe with heavy compression
    start_search = max(0, cut_timestamp - 300) 
    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'{start_search}%{cut_timestamp}', input_video
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    frames = json.loads(result.stdout).get('frames', [])
    
    best_keyframe = 0.0 # Default to the start of the video if no keyframe is found
    
    for frame in frames:
        # Safely try multiple timestamp keys
        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)
            # Find the highest keyframe timestamp that is strictly less than or equal to our cut point
            if keyframe_time <= cut_timestamp and keyframe_time > best_keyframe:
                best_keyframe = keyframe_time
                
    return best_keyframe

def smart_cut_tail(input_video, output_video, tail_duration):
    """Executes the Reverse Smart Cut process: Copy the bulk, re-encode the tiny tail tip, and stitch."""
    
    # 1. Determine exactly where to cut
    target_duration = get_video_duration(input_video)
    if not target_duration:
        print(f"Skipping {input_video}: Could not determine its total duration.")
        return
        
    cut_timestamp = target_duration - tail_duration
    
    if cut_timestamp <= 0:
        print(f"Skipping {input_video}: Video is shorter than the tail reference clip.")
        return

    print(f"Processing: {input_video} (Cutting tail at {cut_timestamp:.2f}s)...")
    
    # 2. Find the preceding keyframe
    prev_keyframe = get_previous_keyframe(input_video, cut_timestamp)

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

    try:
        # 3. Copy the bulk of the video (from 0 to the preceding keyframe)
        if prev_keyframe > 0.0:
            subprocess.run([
                'ffmpeg', '-y', '-v', 'error', '-i', input_video, 
                '-t', str(prev_keyframe), 
                '-c:v', 'copy', '-c:a', 'copy', part1
            ], check=True)
        
        # 4. Re-encode the tiny tip (from the preceding keyframe to the exact cut timestamp)
        tiny_duration = cut_timestamp - prev_keyframe
        subprocess.run([
            'ffmpeg', '-y', '-v', 'error', 
            '-ss', str(prev_keyframe), '-i', input_video, 
            '-t', str(tiny_duration), 
            '-c:v', 'libx264', '-crf', '18', '-c:a', 'copy', part2
        ], check=True)
        
        # 5. Concatenate them using a text list (Forced UTF-8 encoding for emojis/special chars)
        with open(concat_list, 'w', encoding='utf-8') as f:
            if prev_keyframe > 0.0:
                f.write(f"file '{part1}'\n")
            f.write(f"file '{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:
        print(f"Error processing {input_video}. FFmpeg failed.")
    finally:
        # 6. 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():
    if len(sys.argv) < 2:
        print("Usage: python cut_tail.py <path_to_file_or_folder>")
        return
        
    # Get the absolute path and strip any accidental quotes
    raw_input = sys.argv[1].strip(' "''')
    input_arg = os.path.abspath(raw_input)
    
    # Determine if it's a file or a folder
    if os.path.isfile(input_arg):
        if not input_arg.lower().endswith('.mp4'):
            print("Error: The specified file is not an .mp4 video.")
            return
        target_dir = os.path.dirname(input_arg)
        target_files = [os.path.basename(input_arg)]
        
    elif os.path.isdir(input_arg):
        target_dir = input_arg
        target_files = None
        
    else:
        print(f"Error: Path '{input_arg}' does not exist.")
        return

    # Change the working directory to the target folder
    os.chdir(target_dir)

    # If it was a folder, grab all .mp4 files
    if target_files is None:
        target_files = glob.glob("*.mp4")

    # WE NOW LOOK FOR 'tail_ref.mp4' INSTEAD OF 'baseline.mp4'
    tail_ref_file = "o.mp4"

    # Check if tail reference exists in the target directory
    if not os.path.exists(tail_ref_file):
        print(f"Error: '{tail_ref_file}' not found in the target folder: {target_dir}")
        return

    # Remove the reference file from the processing list so we don't try to cut it
    if tail_ref_file in target_files:
        target_files.remove(tail_ref_file)

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

    # Get exact duration of the tail reference
    print(f"Analyzing tail reference duration: {tail_ref_file}...")
    tail_duration = get_video_duration(tail_ref_file)
    
    if not tail_duration:
        print("Error: Could not determine the duration of the tail reference video.")
        return
        
    print(f"Tail reference duration found: {tail_duration} seconds.")

    # Create output directory
    output_dir = "processed_videos"
    os.makedirs(output_dir, exist_ok=True)

    print(f"\nFound {len(target_files)} video(s) to process in '{target_dir}'.\n" + "-"*30)

    # Process each video
    for video in target_files:
        output_path = os.path.join(output_dir, video)
        smart_cut_tail(video, output_path, tail_duration)

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

if __name__ == "__main__":
    main()