import os
import sys
import subprocess
import re

def find_matching_subtitles(video_path, search_dir):
    base_name = os.path.splitext(os.path.basename(video_path))[0]
    subtitle_files = []
    for file in os.listdir(search_dir):
        if re.match(rf"^{re.escape(base_name)}__.*\.srt$", file):
            subtitle_files.append(os.path.join(search_dir, file))
    return sorted(subtitle_files)

def embed_subtitles(video_path, subtitle_paths, output_path):
    try:
        cmd = ["ffmpeg", "-i", video_path]

        for subtitle in subtitle_paths:
            cmd += ["-i", subtitle]

        cmd += ["-c:v", "copy", "-c:a", "copy"]

        if output_path.lower().endswith(".mp4"):
            cmd += ["-c:s", "mov_text"]
        else:
            cmd += ["-c:s", "copy"]

        cmd += ["-map", "0"]
        for i in range(1, len(subtitle_paths) + 1):
            cmd += ["-map", str(i)]

        cmd += ["-y", output_path]

        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode != 0:
            print(f"Error embedding subtitles into {video_path}:\n{result.stderr}")
        else:
            print(f"Embedded subtitles into: {output_path}")

    except Exception as e:
        print(f"Failed to embed subtitles: {e}")

def process_path(input_path):
    if os.path.isfile(input_path):
        directory = os.path.dirname(input_path)
        subtitles = find_matching_subtitles(input_path, directory)
        if subtitles:
            output_file = os.path.splitext(input_path)[0] + "_subbed" + os.path.splitext(input_path)[1]
            embed_subtitles(input_path, subtitles, output_file)
        else:
            print(f"No matching subtitles found for: {input_path}")
    elif os.path.isdir(input_path):
        for file_name in os.listdir(input_path):
            full_path = os.path.join(input_path, file_name)
            if os.path.isfile(full_path) and file_name.lower().endswith(('.mp4', '.mkv', '.avi', '.mov')):
                process_path(full_path)
    else:
        print(f"Invalid input path: {input_path}")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python auto_embed_subtitles.py <video_file_or_directory>")
        sys.exit(1)

    input_path = sys.argv[1]
    process_path(input_path)