"""Traverse a all python files from the route directory gathering all of the
imports. This is useful for putting together requirements files.
"""
from collections import namedtuple
import ast
import argparse
import pathlib
import os
# import pprint as pp


_PROG_NAME = 'import-requirements'
"""The program name (`str`)
"""
_DESC = __doc__
"""The program description (`str`)
"""
# Idea taken from: https://stackoverflow.com/questions/9008451
Import = namedtuple("Import", ["module", "name", "alias", "path"])
"""Will hold all the imports (`namedtuple`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script.
    """
    # Initialise and parse the command line arguments
    parser = _init_cmd_args()
    args = _parse_cmd_args(parser)
    imports = list()
    for d, sd, f in os.walk(args.root_dir):
        if d.endswith('__pycache__'):
            continue
        for f in pathlib.Path(d).glob("*.py"):
            f = pathlib.PurePath(f)
            try:
                rel_path = f.relative_to(args.root_dir)
            except AttributeError:
                rel_path = f.relative_to(args.root_dir)
            imports.extend(
                [(i.module + i.name, rel_path) for i in get_imports(f)]
            )
    if args.full is False:
        imports = [([i[0]], p) for i, p in imports]

    out_imports = []
    max_len = 0
    for i, p in imports:
        import_path = ".".join(i)
        out_imports.append((import_path, p))
        max_len = max(len(import_path), max_len)

    if args.file is True:
        for i, p in sorted(out_imports, key=lambda x: x[0]):
            print(f"{i.ljust(max_len+5)}{str(p)}")
    else:
        for i in sorted(set([i for i, p in out_imports])):
            print(i)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args():
    """Initialise the command line arguments and return the parser.

    Returns
    -------
    args : `argparse.ArgumentParser`
        The argparse parser object with arguments added.
    """
    parser = argparse.ArgumentParser(
        description=_DESC
    )
    parser.add_argument(
        'root_dir', type=str, help="The repository root directory."
    )
    parser.add_argument(
        '--full', action='store_true',
        help="Show the full import path, if this is not used only top level"
        " name is used i.e. ``from <top> import x`` or ``import <top>``."
    )
    parser.add_argument(
        '--file', action='store_true',
        help="Show the Python file names (relative to the ``root_dir``)."
    )
    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):
    """Parse the command line arguments.

    Parameters
    ----------
    parser : `argparse.ArgumentParser`
        The argparse parser object with arguments added.

    Returns
    -------
    args : `argparse.Namespace`
        The argparse namespace object containing the arguments.
    """
    args = parser.parse_args()
    args.root_dir = os.path.expanduser(args.root_dir)
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_imports(path):
    """Get all the import/from x import y statements from the Python file in
    path.

    Parameters
    ----------
    path : `str` or `pathlib.Path`
        The path to the Python script.

    Returns
    -------
    all_imports : `list` of `Import`
        Named tuples representing the imports.
    """
    all_imports = []
    with open(path) as fh:
        root = ast.parse(fh.read(), path)

    for node in ast.iter_child_nodes(root):
        if isinstance(node, ast.Import):
            module = []
        elif isinstance(node, ast.ImportFrom):
            try:
                module = node.module.split('.')
            except AttributeError:
                if node.module is None:
                    continue
        else:
            continue

        for n in node.names:
            all_imports.append(
                Import(module, n.name.split('.'), n.asname, path)
            )
    return all_imports


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
if __name__ == '__main__':
    main()
