import os
from pathlib import Path
import subprocess
from typing import Any, Set
import json

# Work with Amazon Prime, TF1+
# Do not work with dinsey+

url = "URL"

filename = "FILE"
license_key = [
    "--key", "XX",
]

crypted_folder = Path("./raw/")
decrypted_folder = Path("./decrypted/")
decrypted_folder.mkdir(exist_ok=True)

result = subprocess.run(
    ["yt-dlp", "--allow-unplayable", "-J", url],
    capture_output=True,
    text=True
)

video_info = json.loads(result.stdout)

available_languages: Set[str] = set()
for f in video_info.get("formats", []):
    if f.get("acodec") != "none" and f.get("format_id"):
        available_languages.add(f.get("format_id").split("=")[0])
        print(f.get("format_id") + " | " + f.get("format_id").split("=")[0])

yt_dlp_filter = "bv+" + "+".join([f"bestaudio[format_id~='{lang}=']" for lang in available_languages])

yt_dlp_cmd = [
    "yt-dlp",
    "--allow-unplayable",
    "--audio-multistreams",
    "-o", str(crypted_folder) + f"/{filename}.%(ext)s",
    "-f", yt_dlp_filter,
    url
]

print("Running command:", " ".join(yt_dlp_cmd))
subprocess.run(yt_dlp_cmd)

crypted_video_file = list(crypted_folder.glob(filename + ".*.mp4"))[0]
crypted_audio_files = list(crypted_folder.glob(filename + ".*.m4a"))

for audio_file in crypted_audio_files:
    output_file = decrypted_folder / f"{audio_file.stem}.final.m4a"
    if os.path.exists(output_file):
        continue 

    cmd = [
        "mp4decrypt",
        *license_key,
        str(audio_file),
        str(output_file)
    ]
    
    print("Running command:", " ".join(cmd))
    subprocess.run(cmd)

if not os.path.exists(str(decrypted_folder / f"{crypted_video_file.stem}.final.mp4")):
    cmd = [
        "mp4decrypt",
        *license_key,
        str(crypted_video_file),
        str(decrypted_folder / f"{crypted_video_file.stem}.final.mp4")
    ]

    print("Running command:", " ".join(cmd))
    subprocess.run(cmd)

decrypted_video_file = decrypted_folder / f"{crypted_video_file.stem}.final.mp4"
decrypted_audio_files = list(decrypted_folder.glob(filename + ".*.m4a"))

final_file = decrypted_folder / (filename + ".final.mkv")

cmd = [
    "ffmpeg",
    "-nostdin",
    "-i", str(decrypted_video_file),
    *[["-i", audio_file] for audio_file in decrypted_audio_files],
    "-map", "0:v",
    *[["-map",  f"{i+1}:a"] for i, _ in enumerate(decrypted_audio_files)],
    "-c:v", "copy", "-c:a", "copy",
    str(final_file),
]

def flatten(lst: list[Any]):
    for el in lst:
        if isinstance(el, list):
            yield from flatten(el)
        else:
            yield str(el)

subprocess.run(flatten(cmd))
print("Running command:", " ".join(flatten(cmd)))


# Clean up