main

from pyds_functions import *



if __name__ == "__main__":

user_choice = "l" # Used in main loop
sizes_dict = {} # Stores directory as key and sum of files in that directory (in bytes) as value
while True:

    if user_choice == "l":
        print()
        print("Welcome to pyDiskSpace")
        print()
        # Get all drives on system
        drives = get_drives()

        # Print drives to screen
        for i, drive in enumerate(drives):
            print(str(i) + ")", drive)
        print()

        # Ask user to select a drive
        while True:
            try:
                drive = int(input("Please select a drive: "))
                os.chdir(drives[drive])
                break

            except ValueError:
                print("Invalid input.")

            except IndexError:
                print("Invalid drive.")


        # Print directories in selected drive to screen
        print_dirs()

        user_choice = None

    elif user_choice == "u":
        os.chdir("..")
        print_dirs()

    elif user_choice == "q":
        break

    elif user_choice == "s":

        if os.getcwd() in  sizes_dict:
            # Lists each folder inside the directory and it's total size
            return_list = []
            for folder in return_dirs(os.getcwd()):
                if os.path.isdir(os.getcwd() + "\\" + folder):
                    if len(os.getcwd()) >= 4:
                        foldersize = folder_size(sizes_dict, os.getcwd() + "\\" + folder)
                    else:
                        foldersize = folder_size(sizes_dict, os.getcwd() + folder)
                    return_list.append((folder, foldersize))

            print()
            print("Current directory: ", os.getcwd())
            print()

            for folder, size in sorted(return_list, key=lambda student: student[1], reverse=True):
                print(folder, formatbytes(size))

            print()
            input("Press any key to continue...")
            print_dirs()
        else:

            # Identify how many folders in total
            progress_goal = 0
            for root, dirs, files in os.walk(os.getcwd()):
                progress_goal += 1

            # Gets the size of all directories, including subdirectories
            current_progress = 0
            for root, dirs, files in os.walk(os.getcwd()):
                current_progress += 1
                progress = current_progress / progress_goal * 100
                print("Current progress", "{0:.2f}%".format(progress))
                if root not in sizes_dict:

                    sizes_dict[root] = dir_getsize(root)

            # Lists each folder inside the directory and it's total size
            return_list = []
            for folder in return_dirs(os.getcwd()):
                if os.path.isdir(os.getcwd() + "\\" + folder):
                    if len(os.getcwd()) >= 4:
                        foldersize = folder_size(sizes_dict, os.getcwd() + "\\" + folder)
                    else:
                        foldersize = folder_size(sizes_dict, os.getcwd() + folder)
                    return_list.append((folder, foldersize))

            print()
            print("Current directory: ", os.getcwd())
            print()

            for folder, size in sorted(return_list, key=lambda student: student[1], reverse=True):
                print(folder, formatbytes(size))

            print()
            input("Press any key to continue...")
            print_dirs()

    # Navigate to directory if user selected a one2
    else:
        try:
            user_choice = int(user_choice)
            folders = return_dirs(os.getcwd())
            os.chdir(os.getcwd() + "\\" + folders[user_choice])

            # Print directory in new directory to screen
            print_dirs()

        except TypeError:
            pass
        except IndexError:
            print("Directory does not exist.")
        except ValueError:
            print("Invalid input.")

    user_choice = input("Select number above, (u)p a directory, (s)elect this directory, (l)ist drives or (q)uit: ")