import os
from collections import defaultdict


def get_files(dir_name):
    """Gets all the files from the specified directory

    Starting from the specified directory, gathers all the files in the
    defaultdict (dictionary-like object), list of files of particular type
    can be pulled out using ".extension" key on the resulting dictionary.

        files = get_files(my_directory)
        txt_files = files[".txt"]
        exe_files = files[".exe"]

    Input value dir_name can be either a name of a directory placed in the
    current working directory (os.getcwd()) or it's full path.
    """

    if os.path.exists(dir_name) and os.path.isdir(dir_name):
        dirs = [dir_name]
        files = defaultdict(list)
        while dirs:
            current_dir = dirs.pop()
            for f in os.listdir(current_dir):
                if os.path.isdir(os.path.join(current_dir, f)):
                    dirs.append(os.path.join(current_dir, f))
                else:
                    ext = os.path.splitext(f)[1]
                    files[ext].append(f)
        return files
    raise NotADirectoryError(
            "Given directory name doesn't refer to a valid directory."
    )


def print_file(files, extension):
    """Takes a list of files and their extension, prints in nice format."""
    if files:
        print("{} files:\n".format(extension))
        print("\n".join("- {}".format(f) for f in files))
        print("\n")
    else:
        print("no files of type {} found\n".format(extension))


def print_files(mapping, extensions=None):
    """Prints files of specified extensions (all by default).

    Prints the files of specified extensions (all by default) based on
    the mapping received from the get_files function.
    """
    if extensions is not None:
        for ext in extensions:
            print_file(mapping["." + ext], ext)
    else:
        for ext in mapping:
            print_file(mapping[ext], ext.lstrip('.'))


def main():
    import argparse
    parser = argparse.ArgumentParser(description="get files from a directory")
    parser.add_argument('dir_name', help="name of a directory (or full path)")
    parser.add_argument('-t', help="file type(s) to output", nargs='+',
                        dest='extensions', metavar='ext')
    args = parser.parse_args()
    try:
        mapping = get_files(args.dir_name)
        print('\n')
        print_files(mapping, args.extensions)
    except NotADirectoryError as e:
        print(e)


if __name__ == '__main__':
    main()