from pathlib import Path
from tinytag import TinyTag

select_extensions = (".mp3", ".m4a")

def main():
    current_path = Path("./")
    music_files = []

    for e in select_extensions:
        music_files.extend(current_path.rglob(f"*{e}"))

    print(f"Discovered {len(music_files)} music files.")
    yn = input("Do you want to sort them? (y or n) ")

    if yn == "y":
        for file in music_files:
            tags = TinyTag.get(file)
            current_file = Path(file)
            new_path = current_path / f"{tags.albumartist or tags.artist}" / f"{tags.album}" / current_file.name
            new_path.parent.mkdir(parents=True, exist_ok=True)
            current_file.rename(new_path)
            print(f"Moved {current_file} to {new_path}.")
        print(f"Succeeded in sorting {len(music_files)} music files.")

if __name__ == "__main__":
    main()