import os
import glob
import subprocess
import sys
import langcodes
import pycountry
from ass_editor import ASS_Editor, attach_font


def find_files(folder_path):
    """Find all video (MP4, TS, MKV), AAC, and subtitle files in the folder."""
    video_files = (glob.glob(os.path.join(folder_path, "*.mkv")) +
                   glob.glob(os.path.join(folder_path, "*.mp4")) +
                   glob.glob(os.path.join(folder_path, "*.ts")))
    base_video_files = [f for f in video_files if not f.endswith("_DUB.mkv") and
                        not f.endswith("_DUB.mp4") and not f.endswith("_DUB.ts")]
    dub_video_files = [f for f in video_files if f.endswith("_DUB.mkv") or
                       f.endswith("_DUB.mp4") or f.endswith("_DUB.ts")]
    aac_files = (glob.glob(os.path.join(folder_path, "*.aac")) +
                 glob.glob(os.path.join(folder_path, "*.m4a")) +
                 glob.glob(os.path.join(folder_path, "*.mka")))
    subtitle_files = (glob.glob(os.path.join(folder_path, "*.srt")) +
                      glob.glob(os.path.join(folder_path, "*.ass")) +
                      glob.glob(os.path.join(folder_path, "*.vtt")))
    return base_video_files, dub_video_files, aac_files, subtitle_files

def get_base_name(file_path):
    """Extract base name by removing '_DUB' and file extension for video files."""
    file_name = os.path.basename(file_path)
    for ext in ["_DUB.mkv", "_DUB.mp4", "_DUB.ts"]:
        if file_name.endswith(ext):
            return file_name[:-len(ext)]
    return os.path.splitext(file_name)[0]

def get_lang_code(file_path):
    """Helper to extract language code from filename."""
    parts = os.path.splitext(os.path.basename(file_path))[0].split('.')
    lang = parts[-1]
    if len(lang) == 2:
        try:
            return langcodes.Language.make(lang).to_alpha3()
        except:
            return "und"
    elif len(lang) == 3:
        return lang
    return "und"

def group_files(video_files, dub_video_files, aac_files, subtitle_files):
    """Group video, AAC, and subtitle files by base name."""
    file_groups = {}
    for video_file in video_files:
        base_name = get_base_name(video_file)
        file_groups[base_name] = {'video': video_file, 'dub_video': None, 'audio': [], 'subtitles': []}

    for dub_video_file in dub_video_files:
        base_name = get_base_name(dub_video_file)
        if base_name in file_groups:
            file_groups[base_name]['dub_video'] = dub_video_file

    for aac_file in aac_files:
        base_name = os.path.splitext(os.path.basename(aac_file))[0].split('.')[0]
        lang = get_lang_code(aac_file)
        if base_name in file_groups:
            file_groups[base_name]['audio'].append((aac_file, lang))

    for sub_file in subtitle_files:
        base_name = os.path.splitext(os.path.basename(sub_file))[0].split('.')[0]
        lang = get_lang_code(sub_file)
        if base_name in file_groups:
            file_groups[base_name]['subtitles'].append((sub_file, lang))
    return file_groups

def embed_files(folder_path):
    """Embed audio and subtitles using mkvmerge."""
    output_folder = os.path.join(folder_path, "Output")
    os.makedirs(output_folder, exist_ok=True)
    video_files, dub_video_files, aac_files, subtitle_files = find_files(folder_path)
    
    file_groups = group_files(video_files, dub_video_files, aac_files, subtitle_files)

    for base_name, files in file_groups.items():
        video_file = files['video']
        dub_video_file = files['dub_video']
        audio_inputs = files['audio']
        subtitle_inputs = files['subtitles']
        
        output_file = os.path.join(output_folder, base_name + ".mkv")

        if not video_file:
            continue

        # Start mkvmerge command
        cmd = ["mkvmerge", "-o", output_file]

        # 1. Base Video File
        # Set a title for the first video track if desired
        cmd.extend(["--track-name", "0:SeFree", video_file])

        # 2. Dubbed Video File (if exists, add its audio/subs)
        if dub_video_file:
            # Assume track 0 is video (skipped), track 1 is Japanese, track 2 is Thai
            # This logic depends on the internal track IDs of the dub file
            cmd.extend([
                "--no-video", 
                "--language", "2:jpn", "--track-name", "2:Japanese",
                "--language", "3:tha", "--track-name", "3:Thai",
                dub_video_file
            ])

        # 3. External Audio Files
        if not dub_video_file:
            for aac_file, lang in audio_inputs:
                # mkvmerge track ID for external single-track files is usually 0
                cmd.extend([
                    "--language", f"0:{lang}",
                    "--track-name", f"0:{pycountry.languages.get(alpha_3=lang).name if lang != 'und' else 'Unknown'}",
                    aac_file
                ])

        # 4. External Subtitle Files
        for sub_file, lang in subtitle_inputs:
            cmd.extend([
                "--language", f"0:{lang}",
                "--track-name", f"0:{pycountry.languages.get(alpha_3=lang).name if lang != 'und' else 'Unknown'}_BLBL",
                sub_file
            ])
            # cmd=attach_font(cmd, "BLBL")

        print(f"Processing {base_name} with mkvmerge...")
        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            print(f"✅ Successfully created {os.path.basename(output_file)}")
        except subprocess.CalledProcessError as e:
            print(f"❌ Error processing {base_name}: {e.stderr}")

def main():
    os.umask(0o000)
    if len(sys.argv) != 2:
        print("Usage: python add_subtitles_to_mkv.py <folder_path>")
        sys.exit(1)

    folder_path = sys.argv[1].strip()
    if not os.path.exists(folder_path):
        print(f"Error: Folder '{folder_path}' does not exist!")
        sys.exit(1)

    embed_files(folder_path)
    print("🎉 Processing complete!")

if __name__ == "__main__":
    main()