#!/usr/bin/env python3

import json
import subprocess
import sys
from pathlib import Path
import shutil

def run(cmd):
    print(f"[*] Running: {' '.join(map(str, cmd))}")
    subprocess.run(cmd, check=True)

def get_tracks(video):
    result = subprocess.run(
        ["mkvmerge", "-J", str(video)],
        capture_output=True,
        text=True,
        check=True,
    )
    data = json.loads(result.stdout)
    return [t for t in data["tracks"] if t["type"] == "audio"]

def codec_to_ext(codec):
    codec = codec.lower()
    if "aac" in codec: return "aac"
    if "ac-3" in codec or "ac3" in codec: return "ac3"
    if "dts" in codec: return "dts"
    if "flac" in codec: return "flac"
    if "opus" in codec: return "opus"
    return "mka"

def clean(text):
    return "".join(c if c.isalnum() or c in "-_." else "_" for c in text)

def extract(vdir, outdir):
    outdir.mkdir(parents=True, exist_ok=True)

    for video in vdir.iterdir():
        if not video.is_file():
            continue

        name = video.stem
        print(f"\n=== Scanning {video.name} ===")

        for track in get_tracks(video):
            tid = track["id"]
            props = track.get("properties", {})
            lang = props.get("language", "und")
            tname = props.get("track_name", "")
            codec = track.get("codec", "")

            ext = codec_to_ext(codec)

            suffix = lang
            if tname:
                suffix += "-" + clean(tname)

            outfile = outdir / f"{name}-{suffix}.{ext}"

            print(f"Extracting track {tid}: lang={lang}, name={tname}, codec={codec}")
            run(["mkvextract", "tracks", str(video), f"{tid}:{outfile}"])

def include(vdir, adir):
    for video in vdir.iterdir():
        if not video.is_file():
            continue

        name = video.stem
        tmp = video.with_name(name + "-tmp.mkv")

        print(f"\n=== Processing {video.name} ===")

        audios = list(adir.glob(f"{name}-*"))
        if not audios:
            print("No matching audio found.")
            continue

        unknown = []
        known = []

        for a in audios:
            meta = a.stem[len(name)+1:]  # remove "name-"
            parts = meta.split("-", 1)
            lang = parts[0]
            tname = parts[1] if len(parts) > 1 else ""

            entry = (a, lang, tname)
            if lang in ("und", "unknown", ""):
                unknown.append(entry)
            else:
                known.append(entry)

        cmd = ["mkvmerge", "-o", str(tmp), str(video)]

        for file, lang, tname in unknown + known:
            print(f"Adding {file.name} lang={lang} name={tname}")
            cmd += ["--language", f"0:{lang}"]
            if tname:
                cmd += ["--track-name", f"0:{tname}"]
            cmd.append(str(file))

        run(cmd)
        shutil.move(tmp, video)
        print(f"Finished {video.name}")

def main():
    if len(sys.argv) < 3:
        print("Usage:")
        print("  script extract <video_dir> [audio_out_dir]")
        print("  script include <video_dir> [audio_dir]")
        sys.exit(1)

    action = sys.argv[1]
    vdir = Path(sys.argv[2]).resolve()
    adir = Path(sys.argv[3]).resolve() if len(sys.argv) > 3 else vdir

    if not vdir.is_dir():
        print("Video directory not found.")
        sys.exit(1)

    if action == "extract":
        extract(vdir, adir)
    elif action == "include":
        include(vdir, adir)
    else:
        print("Unknown action.")
        sys.exit(1)

if __name__ == "__main__":
    main()