import subprocess
import sys
import os
import json

# Function to convert timestamp to seconds
def timestamp_to_seconds(timestamp):
    hours, minutes, seconds = map(float, timestamp.split(":"))
    return hours * 3600 + minutes * 60 + seconds

# Function to convert seconds back to HH:MM:SS format for mkvmerge
def seconds_to_hms(seconds):
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds = seconds % 60
    return f"{hours:02}:{minutes:02}:{seconds:.3f}"


# Function to get chapters from the MKV file using ffprobe
def get_chapters(input_file):
    cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_chapters", input_file]
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print("Error: Unable to retrieve chapter information.")
        sys.exit(1)

    chapters = json.loads(result.stdout)["chapters"]
    chapter_times = [float(chapter["start_time"]) for chapter in chapters]  # Convert to float
    return chapter_times

# Function to find the nearest chapter time to a given timestamp
def find_nearest_chapter(timestamp, chapter_times):
    timestamp_seconds = timestamp_to_seconds(timestamp)
    nearest_chapter = min(chapter_times, key=lambda chapter: abs(chapter - timestamp_seconds))
    return nearest_chapter

# Function to split the video using mkvmerge
# Function to split the video using mkvmerge
def split_video(input_file, start_time, end_time, output_file):
    # Convert start_time and end_time to HH:MM:SS format for mkvmerge
    start_time_hms = seconds_to_hms(start_time)
    end_time_hms = seconds_to_hms(end_time)
    
    cmd = [
        "mkvmerge",
        "-o", output_file,
        "--split", f"parts:{start_time_hms}-{end_time_hms}",
        input_file
    ]
    
    print(f"Running command: {' '.join(cmd)}")  # Debugging: print the command
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"Error splitting video: {result.stderr}")
        sys.exit(1)


# Main function
def main(input_file, timestamps):
    # Create output folder
    output_folder = "split_output"
    os.makedirs(output_folder, exist_ok=True)
    
    # Get chapter times from the video
    print("Retrieving chapter times...")
    chapter_times = get_chapters(input_file)
    print(f"Found chapter times: {chapter_times}")
    
    # Initialize start time
    start_time = 0.0
    output_index = 1

    # Process each timestamp
    for timestamp in timestamps:
        print(f"Processing timestamp: {timestamp}")
        
        # Find the nearest chapter
        nearest_chapter = find_nearest_chapter(timestamp, chapter_times)
        
        # Generate the output filename
        output_filename = os.path.join(output_folder, f"part_{output_index}.mkv")
        
        # Split the video
        print(f"Splitting from {start_time} to {nearest_chapter}...")
        split_video(input_file, start_time, nearest_chapter, output_filename)
        
        # Update start time for the next part
        start_time = nearest_chapter
        output_index += 1

    # Handle the last segment (from last chapter to the end of the video)
    print(f"Handling final segment from {start_time} to the end of the video...")

    # Get the duration of the video
    cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", input_file]
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print("Error: Unable to retrieve video duration.")
        sys.exit(1)
    
    duration = json.loads(result.stdout)["format"]["duration"]
    print(f"Video duration: {duration}")
    
    # Final segment
    output_filename = os.path.join(output_folder, f"part_{output_index}.mkv")
    print(f"Splitting from {start_time} to {duration}...")
    split_video(input_file, start_time, float(duration), output_filename)
    
    print("Splitting complete. Files saved in split_output.")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python split_mkv_by_chapter.py <input_file> <timestamp1> <timestamp2> ...")
        sys.exit(1)

    input_file = sys.argv[1]
    timestamps = sys.argv[2:]

    main(input_file, timestamps)