#!/usr/bin/env python3
"""
Usage:
    grow [FILES...] [-e <list>] [-d <depth>] [-c] [-s] [-p] [-a] [-i]

Options:
    FILES...
        Target files to search.

    -h, --help
        Prints this menu.

    -c, --create
        Creates a nested directory structure. Reads input from STDIN.

    -s, --size
        Displays file sizes.

    -p, --path
        Displays absolute file paths.

    -a, --all
        Searches all files, including dotfiles.

    -e, --exclude <list>
        Excludes directories from the search. Multiple directories are
        comma-separated.

    -d, --depth <depth>
        Sets search depth limit. Default is 20 directories.

    -i, --info
        Prints various information about the provided directory structure.
"""

import glob
import re
import os
import sys
import stat
from docopt import docopt


#------------------------------------------------------------------------------
#
# Create directory
#
#------------------------------------------------------------------------------


stream = ""
stream_rshift = 0

def valid_indent(stream):
    """
    Validates input stream indentation levels

    -1 is returned if the indentation level is not consistent, otherwise the
    indent coefficient is returned.
    """
    i = 0
    m = None
    for name in stream:
        i += 1
        lspace = len(name) - len(name.lstrip())
        # First file cannot be indented
        if i == 1 and lspace != 0:
            return -1
        # Save the indent of the first indented element in the stream
        elif i > 1 and lspace > 0 and m is None:
            m = lspace
            break

    # We have a flat directory structure
    if m is None:
        return 1

    i = 0
    for name in stream:
        i += 1
        lspace = len(name) - len(name.lstrip())
        # Check for alternate indentation levels. Multiply by 2 to allow for
        # single space/tabs indents.
        if (lspace * 2) % (m * 2) != 0:
            return -1
    return m

def lspace(s):
    """
    Returns the number of leading spaces in 's'
    """
    n = 0
    i = 0
    while s[i] == " " or s[i] == "\t":
        n += 1
        i += 1
    return n

def create_tree():
    """
    Reads an AST representation from stdin and calls grow_tree
    """
    global stream
    global stream_rshift

    stream = sys.stdin.read().rstrip().split("\n")

    stream_rshift = valid_indent(stream)

    if stream_rshift == -1:
        print("[ERROR] Detected alternate indentation levels")
        quit(1)

    grow_tree()

def write_tree_node(path):
    """
    Writes 'path' to disk
    """
    if "." in path:
        try:
            open(os.path.abspath(path), "w")
        except Exception as e:
            print("[ERROR]", e)
    else:
        try:
            os.makedirs(os.path.abspath(path))
        except Exception as e:
            print("[ERROR]", e)

def grow_tree():
    """
    Parses the input stream and writes the directory structure to disk
    """
    print("Creating directory structure...")
    write_tree_node(stream[0])
    parents = []
    for i in range(1, len(stream)):
        name_a = stream[i - 1]
        name_b = stream[i]
        level_a = lspace(name_a)
        level_b = lspace(name_b)
        name_b = name_b.strip()
        # Next node is a top level node
        if level_b == 0:
            parents = []
            write_tree_node(name_b)
            continue
        # Next node is one level down
        if level_b > level_a:
            parents.append(name_a.strip())
            write_tree_node(os.path.join(*parents, name_b))
        # Next node is on the same level
        elif level_b == level_a:
            write_tree_node(os.path.join(*parents, name_b))
        # Next node is one or more levels up
        elif level_b < level_a:
            k = int((level_a - level_b) / stream_rshift)
            parents = parents[:-k]
            write_tree_node(os.path.join(*parents, name_b))
    print("Done!")


#------------------------------------------------------------------------------
#
# Search directory
#
#------------------------------------------------------------------------------


global_show_path    = False
global_show_size    = False
global_search_depth = 20

info_max_depth = 0
info_num_files = 0
info_size      = 0

exclude_list = []

search_pattern = None

rshift = 4

def get_files(root):
    """
    Returns the contents of 'root'
    """
    files = []
    for target in search_pattern.format(root).split(" "):
        files.extend(glob.glob(target, recursive = False))
    return files

def usize(bytes):
    """
    Returns the unit size of 'bytes' as a string
    """
    if bytes > 1e+9:
        return "{:.1f} {}".format(bytes / 1e+9, "GiB")
    elif bytes > 1e+6:
        return "{:.1f} {}".format(bytes / 1e+6, "MiB")
    elif bytes > 1e+3:
        return "{:.1f} {}".format(bytes / 1e+3, "KiB")
    else:
        return "{} {}".format(bytes, "Byte")

def traverse_tree(root, depth=0):
    """
    Performs a recursive depth-first search on 'root' and prints its contents
    """
    if depth > global_search_depth:
        return

    isdir = os.path.isdir(root)
    string = None

    # Show absolute path
    if global_show_path:
        string = os.path.abspath(root)
    else:
        string = os.path.basename(root)

    # Show file size
    if global_show_size and not isdir:
        try:
            string = string + " {}".format(usize(os.stat(root).st_size))
        except OSError:
            pass

    if global_show_path:
        print(string + "/" * isdir)
    else:
        print(" " * depth * rshift + string + "/" * isdir)

    if os.path.basename(root) in exclude_list:
        return

    for file in sorted(get_files(root), key=str):
        if isdir:
            traverse_tree(file, depth + 1)

def traverse_tree_info(root, depth=1):
    """
    Performs a recursive depth-first search on 'root' and gathers information about
    the directory structure
    """
    global info_max_depth
    global info_num_files
    global info_size

    info_max_depth += (depth > info_max_depth)

    for file in get_files(root):
        if os.path.basename(file) in exclude_list:
            continue
        if os.path.isdir(file):
            traverse_tree_info(file, depth + 1)
        else:
            info_num_files += 1
            try:
                info_size += os.stat(file).st_size
            except OSError:
                pass

def print_tree_info(files):
    """
    Calls traverse_tree_info and prints the results
    """
    global info_max_depth
    global info_num_files
    global info_size

    for file in files:
        if not os.path.isdir(file):
            continue

        info_max_depth = 0
        info_num_files = 0
        info_size = 0
        traverse_tree_info(file)

        print("{}\n    depth: {}, files: {}, size: {}\n".format(
            os.path.basename(file),
            info_max_depth,
            info_num_files,
            usize(info_size)
        ))

if __name__ == "__main__":
    argv = docopt(__doc__, help=True)

    if argv["--create"]:
        create_tree()
        quit(0)

    global_show_path = argv["--path"]
    global_show_size = argv["--size"]

    if argv["--all"]:
        search_pattern = "{0}/* {0}/.*"
    else:
        search_pattern = "{}/*"

    if argv["--depth"]:
        global_search_depth = int(argv["--depth"])

    if argv["--exclude"]:
        exclude_list = argv["--exclude"].split(",")

    files = argv["FILES"]

    if len(files) == 0:
        files = get_files(".")

    files = [ x.rstrip("/") for x in files ]

    if argv["--info"]:
        print_tree_info(files)
        quit(0)

    for file in sorted(files, key=str):
        traverse_tree(file)
