import argparse
import sys
import os
import shutil

MEDIA_TYPES = { ".mp4", ".mkv", ".m4v", ".avi", ".webm" }
SUBS_TYPES = { ".srt", ".ass", ".vtt" }


print(sys.argv)
if len(sys.argv) != 2: print("Folder was not provided, exiting") ; exit(2)
folder = sys.argv[1]
if folder == ".": folder = os.curdir

subs = []
subs_ext = []
shifted_subs = []
shifted_subs_ext = []
media = []
names = []
shift = 0


for file in os.scandir(folder):
    i = file.name.rfind(".")
    ext = file.name[i: ]
    if ext in MEDIA_TYPES:
        media.append(file.name)
        names.append(file.name[ :i])
    if ext in SUBS_TYPES:
        subs.append(file.name)
        subs_ext.append(ext)

def rename(shift=0):
    for sub, name, ext in zip(shifted_subs, names, shifted_subs_ext):
        print(f"{sub} -> {name}{ext}")
        shutil.move(sub, name + ext)

def rename_ok():
    global shift
    for sub, name, ext in zip(shifted_subs, names, shifted_subs_ext):
            print(f"---{sub}\n{name}{ext}\n")
    answer = input("\nType a number to shift subs index or 'R' to proceed with renaming\n")
    if answer.lower() == "r": return True
    if answer.isnumeric(): shift = int(answer)
    return False


if len(media) > len(subs):
    print("--Found more media than subs\n")
elif len(media) < len(subs):
    print("--Found more subs than media\n")

while (True):
    print("shift -", shift)
    shifted_subs = []
    shifted_subs_ext = []
    for i in range(len(subs)):
        if i < shift: continue
        shifted_subs.append(subs[i])
        shifted_subs_ext.append(subs[i])

    if rename_ok():
        rename(shift)
        break