import subprocess
import json
from pathlib import Path

INPUT_DIR = Path("/Entertainment_1/Downloads/USCK/Series/Hokkaido.Gals.Are.Super.Adorable.2024.S01.1080p.CR.WEB-DL.DUAL.AAC2.0.H.264-[SeFree]")
OUTPUT_DIR = INPUT_DIR / "output"

# remove only forced thai
TARGET_LANGS = ["tha"]
REMOVE_FORCED = True

# track types to apply
TRACK_TYPES = ["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 split_tracks(tracks):
    keep = {
        "video": [],
        "audio": [],
        "subtitles": [],
    }

    removed = []

    for t in tracks:
        track_id = t["id"]
        track_type = t["type"]
        props = t["properties"]

        lang = props.get("language")
        forced = props.get("forced_track", False)

        should_remove = False

        # remove ONLY forced thai subtitles
        if (
            track_type in TRACK_TYPES
            and TARGET_LANGS
            and lang in TARGET_LANGS
            and REMOVE_FORCED
            and forced
        ):
            should_remove = True

        if should_remove:
            removed.append(track_id)
        else:
            if track_type in keep:
                keep[track_type].append(track_id)

    return keep, removed


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

    tracks = get_tracks(mkv_file)
    keep, removed = split_tracks(tracks)

    if not removed:
        print("  No forced Thai tracks")
        continue

    output_file = OUTPUT_DIR / mkv_file.name

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

    if keep["video"]:
        cmd.extend(["--video-tracks", ",".join(map(str, keep["video"]))])

    if keep["audio"]:
        cmd.extend(["--audio-tracks", ",".join(map(str, keep["audio"]))])

    if keep["subtitles"]:
        cmd.extend(["--subtitle-tracks", ",".join(map(str, keep["subtitles"]))])

    cmd.append(str(mkv_file))

    subprocess.run(cmd)

    print(f"  Removed forced Thai tracks: {removed}")

print("Done")