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

def main():
    """
    Finds all encoded segments in the 'segments' directory and uses mkvmerge
    to mux them into a single, final video file.
    """
    parser = argparse.ArgumentParser(
        description="Muxes encoded segments from the 'segments' directory into a final file using mkvmerge.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("--cleanup", action='store_true', help="Remove the 'segments' and 'cuts' directories after completion.")
    args = parser.parse_args()

    segments_dir = "segments"
    cuts_dir = "cuts"

    # --- Pre-flight Checks ---
    if not shutil.which("mkvmerge"):
        print("Error: 'mkvmerge' not found. Please ensure it is installed and in your system's PATH.")
        sys.exit(1)

    if not os.path.isdir(segments_dir):
        print(f"Error: Input directory '{segments_dir}' not found. Please run the encoder script first.")
        sys.exit(1)

    segment_files = sorted([f for f in os.listdir(segments_dir) if f.endswith('.mkv')])
    if not segment_files:
        print(f"Error: No .mkv files found in the '{segments_dir}' directory.")
        sys.exit(1)

    print(f"Found {len(segment_files)} segments to mux.")

    # --- Determine Output Filename ---
    # Infer the base name from the first segment file.
    # e.g., "My Video_segment001.mkv" -> "My Video"
    try:
        base_name = segment_files[0].rsplit('_segment', 1)[0]
        output_filename = f"temp_{base_name}.mkv"
    except IndexError:
        print("Error: Could not determine a base name from the segment files.")
        print("Files should be named like '..._segmentXXX.mkv'.")
        sys.exit(1)
        
    print(f"Output file will be: {output_filename}")

    # --- Build mkvmerge Command ---
    # mkvmerge -o output.mkv segment1.mkv + segment2.mkv + segment3.mkv ...
    command = ['mkvmerge', '-o', output_filename]
    
    # Add the first file
    command.append(os.path.join(segments_dir, segment_files[0]))
    
    # Add the subsequent files with the '+' append operator
    for segment in segment_files[1:]:
        command.append('+')
        command.append(os.path.join(segments_dir, segment))

    # --- Execute Muxing ---
    print("Running mkvmerge...")
    try:
        process = subprocess.run(command, check=True, capture_output=True, text=True)
        print("mkvmerge output:")
        print(process.stdout)
        print(f"\nSuccess! Segments muxed into '{output_filename}'.")
    except subprocess.CalledProcessError as e:
        print("\nError: mkvmerge failed.")
        print(f"Command: {' '.join(e.cmd)}")
        print(f"Return Code: {e.returncode}")
        print(f"Stderr:\n{e.stderr}")
        sys.exit(1)

    # --- Cleanup ---
    if args.cleanup:
        print("\n--- Cleaning up temporary directories... ---")
        try:
            if os.path.isdir(segments_dir):
                shutil.rmtree(segments_dir)
                print(f"Removed '{segments_dir}' directory.")
            if os.path.isdir(cuts_dir):
                shutil.rmtree(cuts_dir)
                print(f"Removed '{cuts_dir}' directory.")
            print("Cleanup complete.")
        except OSError as e:
            print(f"Error during cleanup: {e}")

if __name__ == "__main__":
    main()