import subprocess
import json
from pathlib import Path

INPUT_DIR = Path("/Entertainment_1/Downloads/USCK/Series/Release.that.Witch.2026.S01.1080p.CR.WEB-DL.AAC2.0.H.264-[SeFree]")
OUTPUT_DIR = Path.joinpath(INPUT_DIR,"output")
SHIFT_MS = 28500  # +15 sec (negative = backward)

# choose one mode
TARGET_TRACK_IDS = None        # example: [1,2]
TARGET_LANGS = ["tha","eng"]         # example: ["eng", "jpn", "tha"]

# track types to shift
TRACK_TYPES = ["subtitles"] #  ["audio", "subtitles"]

OUTPUT_DIR.mkdir(exist_ok=True)


def get_tracks(file):
    cmd = ["mkvmerge", "-J", str(file)]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return json.loads(result.stdout)["tracks"]


def select_tracks(tracks):
    selected = []

    for t in tracks:
        if t["type"] not in TRACK_TYPES:
            continue

        track_id = t["id"]
        lang = t["properties"].get("language")

        if TARGET_TRACK_IDS:
            if track_id in TARGET_TRACK_IDS:
                selected.append(track_id)

        elif TARGET_LANGS:
            if lang in TARGET_LANGS:
                selected.append(track_id)

    return selected


for mkv_file in INPUT_DIR.glob("*.mkv"):
    print(f"Processing: {mkv_file.name}")

    tracks = get_tracks(mkv_file)
    selected_ids = select_tracks(tracks)

    if not selected_ids:
        print("  No matching tracks")
        continue

    output_file = OUTPUT_DIR / mkv_file.name

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

    for track_id in selected_ids:
        cmd.extend(["--sync", f"{track_id}:{SHIFT_MS}"])

    cmd.append(str(mkv_file))

    subprocess.run(cmd)

print("Done")