#!/usr/bin/env python3
import subprocess
import sys
import json
import pathlib

VIDEO_EXTS = {".mkv", ".mp4", ".avi", ".mov"}

LANG_CODES = {
    "eng": "English", "en": "English", "english": "English",
    "spa": "Spanish", "es": "Spanish", "spanish": "Spanish",
    "fra": "French", "fre": "French", "fr": "French", "french": "French",
    "deu": "German", "ger": "German", "de": "German", "german": "German",
}

CANONICAL = {
    "English": "eng",
    "Spanish": "spa",
    "French": "fra",
    "German": "deu",
}

def normalize_language(value):
    key = value.lower()
    if key not in LANG_CODES:
        raise ValueError(f"Unknown language: {value}")
    return CANONICAL[LANG_CODES[key]]

def probe(file):
    cmd = ["ffprobe", "-v", "error", "-print_format", "json", "-show_streams", file]
    return json.loads(subprocess.check_output(cmd))

def collect_files(paths):
    files = []
    for p in paths:
        p = pathlib.Path(p)
        if p.is_dir():
            files += [f for f in p.iterdir() if f.suffix.lower() in VIDEO_EXTS]
        elif p.suffix.lower() in VIDEO_EXTS:
            files.append(p)
    return files

def preserve_original(file):
    original = file.with_name(f"{file.stem}_original{file.suffix}")
    if original.exists():
        print(f"Original already preserved: {original.name}")
        return
    file.rename(original)
    print(f"Preserved original as: {original.name}")

def set_audio_lang(file, lang):
    data = probe(file)
    audio_streams = [s for s in data["streams"] if s["codec_type"] == "audio"]

    if not audio_streams:
        print(f"No audio streams in {file.name}")
        return

    current = audio_streams[0].get("tags", {}).get("language")

    if current:
        ans = input(f"{file.name}: audio already '{current}'. Overwrite? [y/N] ")
        if ans.lower() != "y":
            return

    out = file.with_name(f"{file.stem}_dub-{lang}{file.suffix}")

    print(f"{file.name} -> {out.name}")

    cmd = [
        "ffmpeg", "-y",
        "-i", str(file),
        "-map", "0",
        "-c", "copy",
        "-metadata:s:a:0", f"language={lang}",
        str(out)
    ]

    result = subprocess.run(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=True
    )

    if result.returncode != 0:
        print(f"ERROR processing {file.name}")
        print(result.stderr.strip())
        return


    preserve_original(file)


def main():
    if len(sys.argv) == 1:
        targets = ["."]
        lang = input("Language: ")
    elif len(sys.argv) == 2:
        targets = ["."]
        lang = sys.argv[1]
    else:
        targets = sys.argv[1:-1]
        lang = sys.argv[-1]
    subprocess.run(["python", "/home/honney/.bin/tracker.py", "add", "set_audio_lang"])
    lang = normalize_language(lang)

    for f in collect_files(targets):
        set_audio_lang(f, lang)

if __name__ == "__main__":
    
    main()
