import os
import re
import subprocess

def find_matching_files(folder_path):
    # Regex pattern for DJI files (e.g., DJI_20250530075349_0003_D)
    pattern = r"DJI_(\d{14}_\d{4})_D\.(MP4|WAV)"

    # Dictionaries to store files by their unique identifier
    mp4_files = {}
    wav_files = {}

    # Iterate through files in the folder
    for filename in os.listdir(folder_path):
        match = re.match(pattern, filename)
        if match:
            identifier = match.group(1)  # e.g., 20250530075349_0003
            extension = match.group(2).lower()

            if extension == "mp4":
                mp4_files[identifier] = filename
            elif extension == "wav":
                wav_files[identifier] = filename

    return mp4_files, wav_files

def merge_files(folder_path, mp4_files, wav_files):
    # Ensure output directory exists
    output_dir = os.path.join(folder_path, "merged_output")
    os.makedirs(output_dir, exist_ok=True)

    # Find matching pairs
    for identifier in mp4_files:
        if identifier in wav_files:
            mp4_path = os.path.join(folder_path, mp4_files[identifier])
            wav_path = os.path.join(folder_path, wav_files[identifier])
            output_filename = f"DJI_{identifier}_D_merged.MP4"
            output_path = os.path.join(output_dir, output_filename)

            # FFmpeg command to keep original AAC stereo and add WAV stereo as a second stereo track, excluding timecode
            ffmpeg_cmd = [
                "ffmpeg",
                "-i", mp4_path,           # Input video file with original stereo audio
                "-i", wav_path,           # Input WAV file with stereo audio
                "-c:v", "copy",           # Copy video stream without re-encoding
                "-map", "0:v:0",          # Map video from MP4
                "-map", "0:a:0",          # Map AAC stereo audio from MP4 (first audio stream, index 0)
                "-c:a:0", "copy",         # Copy MP4 AAC audio without re-encoding
                "-map", "1:a:0",          # Map WAV stereo audio
                "-c:a:1", "aac",          # Encode WAV audio to AAC
                "-b:a:1", "192k",         # Set a reasonable bitrate for AAC
                "-ac:a:1", "2",           # Explicitly set WAV audio to stereo (2 channels)
                "-metadata:s:a:0", "title=MP4 Stereo Audio",  # Label first audio track
                "-metadata:s:a:1", "title=WAV Stereo Audio",  # Label second audio track
                "-map_chapters", "-1",    # Exclude chapters (if any)
                "-map_metadata", "-1",    # Exclude global metadata (to avoid timecode inclusion)
                "-shortest",              # Match duration to shortest input
                output_path
            ]

            try:
                # Run FFmpeg command
                result = subprocess.run(ffmpeg_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                print(f"Successfully merged: {output_filename}")
            except subprocess.CalledProcessError as e:
                print(f"Error merging {mp4_files[identifier]} and {wav_files[identifier]}: {e.stderr}")

def main():
    # Get the folder path
    folder_path = input("Enter the folder path containing MP4 and WAV files (or press Enter for current directory): ").strip()
    if not folder_path:
        folder_path = os.getcwd()

    # Validate folder path
    if not os.path.isdir(folder_path):
        print("Error: Invalid folder path.")
        return

    # Find matching MP4 and WAV files
    mp4_files, wav_files = find_matching_files(folder_path)

    if not mp4_files or not wav_files:
        print("No matching MP4 or WAV files found.")
        return

    # Merge matching files
    merge_files(folder_path, mp4_files, wav_files)

if __name__ == "__main__":
    main()