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

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

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",
        str(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.extend(
                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():
        return
    file.rename(original)
    print(f"preserve: {file.name} -> {original.name}")

def set_sub_lang(file, lang):
    data = probe(file)
    subs = [s for s in data["streams"] if s["codec_type"] == "subtitle"]

    if not subs:
        return

    for i, s in enumerate(subs):
        cur = s.get("tags", {}).get("language", "unset")
        print(f"{i}: subtitle ({cur})")

    sel = input(f"{file.name}: select subtitle index (blank = skip): ")
    if sel == "":
        return

    idx = int(sel)
    current = subs[idx].get("tags", {}).get("language")

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

    out = file.with_name(f"{file.stem}_sub-{lang}{file.suffix}")
    print(f"{file.name} -> {out.name}")

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

    result = subprocess.run(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=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("Subtitle language: ")
    elif len(sys.argv) == 2:
        targets = ["."]
        lang = sys.argv[1]
    else:
        targets = sys.argv[1:-1]
        lang = sys.argv[-1]

    lang = normalize_language(lang)

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

if __name__ == "__main__":
    main()