#!/usr/bin/env python3
import subprocess
import json
import os
import sys
import argparse

def get_video_duration(video_path):
    """Gets the duration of a video file in seconds using ffprobe."""
    command = [
        'ffprobe',
        '-v', 'error',
        '-show_entries', 'format=duration',
        '-of', 'default=noprint_wrappers=1:nokey=1',
        video_path
    ]
    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        return float(result.stdout.strip())
    except FileNotFoundError:
        print("\nError: 'ffprobe' command not found. Is it installed and in your system's PATH?")
        sys.exit(1)
    except (subprocess.CalledProcessError, ValueError) as e:
        print(f"\nError getting video duration: {e}")
        sys.exit(1)

def get_best_hwaccel():
    """
    Checks for available FFmpeg hardware acceleration methods and returns the best one.
    Priority is defined as: cuda > qsv > dxva2.
    Returns the name of the best available method, or None if none are found.
    """
    priority = ['cuda', 'qsv', 'dxva2']
    print("Checking for available hardware acceleration...")
    try:
        # Run ffmpeg to list available hardware acceleration methods
        result = subprocess.run(
            ['ffmpeg', '-hide_banner', '-hwaccels'],
            capture_output=True, text=True, encoding='utf-8', check=False
        )
        
        if result.returncode != 0:
            print("Warning: Could not query FFmpeg for hardware acceleration. Using software decoding.")
            return None

        # Parse the output to get a list of available methods
        available_methods = result.stdout.strip().split('\n')
        if len(available_methods) > 1:
            available_methods = available_methods[1:]

        # Check for our preferred methods in order of priority
        for method in priority:
            if method in available_methods:
                print(f"Found best available hardware acceleration: {method}")
                return method
        
        print("No high-priority hardware acceleration (cuda, qsv, dxva2) found. Using software decoding.")
        return None

    except FileNotFoundError:
        # This will be handled by the main ffmpeg call later, so we just return None.
        return None
    except Exception as e:
        print(f"Warning: An error occurred while checking for hwaccel: {e}. Using software decoding.")
        return None

def cut_video_into_scenes(video_path, max_segment_length, hwaccel=None):
    """
    Cuts a video into segments based on a .scenes.json file using FFmpeg's segment muxer,
    ensuring no segment exceeds a maximum length.

    Args:
        video_path (str): The path to the input video file.
        max_segment_length (int): The maximum allowed length for any segment in seconds.
        hwaccel (str, optional): The hardware acceleration method to use for decoding.
    """
    print(f"\nProcessing video: {os.path.basename(video_path)}")

    # 1. Derive the .scenes.json path and check for its existence
    base_name = os.path.splitext(video_path)[0]
    json_path = f"{base_name}.scenes.json"

    if not os.path.isfile(json_path):
        print(f"\nError: Scene file not found at '{json_path}'")
        print("Please run scene_detector.py on the video first.")
        sys.exit(1)

    # 2. Load the scene timestamps from the JSON file
    try:
        with open(json_path, 'r') as f:
            scene_timestamps = json.load(f)
    except json.JSONDecodeError:
        print(f"\nError: Could not parse the JSON file at '{json_path}'. It may be corrupted.")
        sys.exit(1)

    if not scene_timestamps:
        print("Warning: The scene file is empty. No segments to cut.")
        return

    # Get total video duration to handle the last segment correctly
    print("Getting video duration...")
    video_duration = get_video_duration(video_path)
    print(f"Video duration: {video_duration:.2f} seconds.")

    # 3. Create the output directory
    output_dir = "cuts"
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output will be saved to the '{output_dir}' directory.")

    if hwaccel:
        print(f"Using hardware acceleration for decoding: {hwaccel}")

    # 4. Process timestamps to enforce max segment length
    print(f"Enforcing maximum segment length of {max_segment_length} seconds...")
    
    # Define all segment boundaries, starting at 0 and ending at the video's duration
    segment_boundaries = [0.0] + sorted(list(set(scene_timestamps))) + [video_duration]
    
    final_cut_points = []
    for i in range(len(segment_boundaries) - 1):
        start = segment_boundaries[i]
        end = segment_boundaries[i+1]
        
        current_time = start
        while (end - current_time) > max_segment_length:
            current_time += max_segment_length
            final_cut_points.append(current_time)
        
        final_cut_points.append(end)

    # Clean up the final list: sort, remove duplicates, and remove the final cut at the very end
    final_cut_points = sorted(list(set(final_cut_points)))
    if final_cut_points and final_cut_points[-1] >= video_duration - 0.1: # Use a small tolerance
        final_cut_points.pop()

    print(f"Original scenes: {len(scene_timestamps)}. Total segments after splitting: {len(final_cut_points) + 1}.")

    # 5. Prepare arguments for the segment muxer
    segment_times_str = ",".join(map(str, final_cut_points))
    video_basename = os.path.basename(base_name)
    output_pattern = os.path.join(output_dir, f"{video_basename}_segment%03d.mkv")

    # 6. Build the single FFmpeg command
    command = ['ffmpeg', '-hide_banner']
    if hwaccel:
        command.extend(['-hwaccel', hwaccel])
    
    command.extend([
        '-i', video_path,
        '-c:v', 'utvideo',
        '-an', '-sn', '-dn',
        '-map_metadata', '-1',
        '-map_chapters', '-1',
        '-f', 'segment',
        '-segment_times', segment_times_str,
        '-segment_start_number', '1',
        '-reset_timestamps', '1',
        output_pattern
    ])

    print("\nStarting FFmpeg to cut all segments in a single pass...")
    try:
        # Use Popen to show FFmpeg's progress in real-time
        process = subprocess.Popen(command, stderr=subprocess.PIPE, text=True, encoding='utf-8')

        for line in iter(process.stderr.readline, ''):
            if line.strip().startswith('frame='):
                sys.stderr.write(f'\r{line.strip()}')
                sys.stderr.flush()
        
        process.wait()
        sys.stderr.write('\n')
        sys.stderr.flush()

        if process.returncode != 0:
            print(f"\nWarning: FFmpeg may have encountered an error (exit code: {process.returncode}).")
        else:
            total_segments = len(final_cut_points) + 1
            print(f"\nSuccessfully created {total_segments} segments in the '{output_dir}' directory.")

    except FileNotFoundError:
        print("\nError: 'ffmpeg' command not found. Is it installed and in your system's PATH?")
        sys.exit(1)
    except Exception as e:
        print(f"\nAn error occurred during cutting: {e}")

def main():
    """
    Main function to handle command-line arguments.
    """
    parser = argparse.ArgumentParser(
        description="Automatically detects the best hardware decoder and cuts a video into scene-based segments using a corresponding '.scenes.json' file."
    )
    parser.add_argument("video_path", help="Path to the video file.")
    parser.add_argument(
        '--segment-length',
        type=int,
        default=10,
        help='Maximum length of a segment in seconds. Default: 10'
    )
    
    args = parser.parse_args()

    if not os.path.isfile(args.video_path):
        print(f"Error: Input file not found: '{args.video_path}'")
        sys.exit(1)
    
    # Automatically determine the best hardware acceleration method
    hwaccel_method = get_best_hwaccel()
    
    cut_video_into_scenes(args.video_path, args.segment_length, hwaccel_method)

if __name__ == "__main__":
    main()