#!/usr/bin/env python

# Usage (the graphviz package must be installed in your distribution)
#  ./support/scripts/graph-depends [-p package-name] > test.dot
#  dot -Tpdf test.dot -o test.pdf
#
# With no arguments, graph-depends will draw a complete graph of
# dependencies for the current configuration.
# If '-p <package-name>' is specified, graph-depends will draw a graph
# of dependencies for the given package name.
# If '-d <depth>' is specified, graph-depends will limit the depth of
# the dependency graph to 'depth' levels.
#
# Limitations
#
#  * Some packages have dependencies that depend on the Buildroot
#    configuration. For example, many packages have a dependency on
#    openssl if openssl has been enabled. This tool will graph the
#    dependencies as they are with the current Buildroot
#    configuration.
#
# Copyright (C) 2010-2013 Thomas Petazzoni <thomas.petazzoni@free-electrons.com>

import logging
import sys
import subprocess
import argparse
from fnmatch import fnmatch

import brpkgutil

# Modes of operation:
MODE_FULL = 1   # draw full dependency graph for all selected packages
MODE_PKG = 2    # draw dependency graph for a given package

allpkgs = []


# Execute the "make show-targets" command to get the list of the main
# Buildroot PACKAGES and return it formatted as a Python list. This
# list is used as the starting point for full dependency graphs
def get_targets():
    logging.info("Getting targets")
    cmd = ["make", "-s", "--no-print-directory", "show-targets"]
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
    output = p.communicate()[0].strip()
    if p.returncode != 0:
        return None
    if output == '':
        return []
    return output.split(' ')


# Recursive function that builds the tree of dependencies for a given
# list of packages. The dependencies are built in a list called
# 'dependencies', which contains tuples of the form (pkg1 ->
# pkg2_on_which_pkg1_depends, pkg3 -> pkg4_on_which_pkg3_depends) and
# the function finally returns this list.
def get_all_depends(pkgs, get_depends_func):
    dependencies = []

    # Filter the packages for which we already have the dependencies
    filtered_pkgs = []
    for pkg in pkgs:
        if pkg in allpkgs:
            continue
        filtered_pkgs.append(pkg)
        allpkgs.append(pkg)

    if len(filtered_pkgs) == 0:
        return []

    depends = get_depends_func(filtered_pkgs)

    deps = set()
    for pkg in filtered_pkgs:
        pkg_deps = depends[pkg]

        # This package has no dependency.
        if pkg_deps == []:
            continue

        # Add dependencies to the list of dependencies
        for dep in pkg_deps:
            dependencies.append((pkg, dep))
            deps.add(dep)

    if len(deps) != 0:
        newdeps = get_all_depends(deps, get_depends_func)
        if newdeps is not None:
            dependencies += newdeps

    return dependencies


# The Graphviz "dot" utility doesn't like dashes in node names. So for
# node names, we strip all dashes. Also, nodes can't start with a number,
# so we prepend an underscore.
def pkg_node_name(pkg):
    return "_" + pkg.replace("-", "")


# Basic cache for the results of the is_dep() function, in order to
# optimize the execution time. The cache is a dict of dict of boolean
# values. The key to the primary dict is "pkg", and the key of the
# sub-dicts is "pkg2".
is_dep_cache = {}


def is_dep_cache_insert(pkg, pkg2, val):
    try:
        is_dep_cache[pkg].update({pkg2: val})
    except KeyError:
        is_dep_cache[pkg] = {pkg2: val}


# Retrieves from the cache whether pkg2 is a transitive dependency
# of pkg.
# Note: raises a KeyError exception if the dependency is not known.
def is_dep_cache_lookup(pkg, pkg2):
    return is_dep_cache[pkg][pkg2]


# This function return True if pkg is a dependency (direct or
# transitive) of pkg2, dependencies being listed in the deps
# dictionary. Returns False otherwise.
# This is the un-cached version.
def is_dep_uncached(pkg, pkg2, deps):
    try:
        for p in deps[pkg2]:
            if pkg == p:
                return True
            if is_dep(pkg, p, deps):
                return True
    except KeyError:
        pass
    return False


# See is_dep_uncached() above; this is the cached version.
def is_dep(pkg, pkg2, deps):
    try:
        return is_dep_cache_lookup(pkg, pkg2)
    except KeyError:
        val = is_dep_uncached(pkg, pkg2, deps)
        is_dep_cache_insert(pkg, pkg2, val)
        return val


# This function eliminates transitive dependencies; for example, given
# these dependency chain: A->{B,C} and B->{C}, the A->{C} dependency is
# already covered by B->{C}, so C is a transitive dependency of A, via B.
# The functions does:
#   - for each dependency d[i] of the package pkg
#     - if d[i] is a dependency of any of the other dependencies d[j]
#       - do not keep d[i]
#     - otherwise keep d[i]
def remove_transitive_deps(pkg, deps):
    d = deps[pkg]
    new_d = []
    for i in range(len(d)):
        keep_me = True
        for j in range(len(d)):
            if j == i:
                continue
            if is_dep(d[i], d[j], deps):
                keep_me = False
        if keep_me:
            new_d.append(d[i])
    return new_d


# List of dependencies that all/many packages have, and that we want
# to trim when generating the dependency graph.
MANDATORY_DEPS = ['toolchain', 'skeleton']


# This function removes the dependency on some 'mandatory' package, like the
# 'toolchain' package, or the 'skeleton' package
def remove_mandatory_deps(pkg, deps):
    return [p for p in deps[pkg] if p not in MANDATORY_DEPS]


# This function will check that there is no loop in the dependency chain
# As a side effect, it builds up the dependency cache.
def check_circular_deps(deps):
    def recurse(pkg):
        if pkg not in list(deps.keys()):
            return
        if pkg in not_loop:
            return
        not_loop.append(pkg)
        chain.append(pkg)
        for p in deps[pkg]:
            if p in chain:
                logging.warning("\nRecursion detected for  : %s" % (p))
                while True:
                    _p = chain.pop()
                    logging.warning("which is a dependency of: %s" % (_p))
                    if p == _p:
                        sys.exit(1)
            recurse(p)
        chain.pop()

    not_loop = []
    chain = []
    for pkg in list(deps.keys()):
        recurse(pkg)


# This functions trims down the dependency list of all packages.
# It applies in sequence all the dependency-elimination methods.
def remove_extra_deps(deps, rootpkg, transitive):
    for pkg in list(deps.keys()):
        if not pkg == rootpkg:
            deps[pkg] = remove_mandatory_deps(pkg, deps)
    for pkg in list(deps.keys()):
        if not transitive or pkg == rootpkg:
            deps[pkg] = remove_transitive_deps(pkg, deps)
    return deps


# Print the attributes of a node: label and fill-color
def print_attrs(outfile, pkg, version, depth, colors):
    name = pkg_node_name(pkg)
    if pkg == 'all':
        label = 'ALL'
    else:
        label = pkg
    if depth == 0:
        color = colors[0]
    else:
        if pkg.startswith('host') \
                or pkg.startswith('toolchain') \
                or pkg.startswith('rootfs'):
            color = colors[2]
        else:
            color = colors[1]
    if version == "virtual":
        outfile.write("%s [label = <<I>%s</I>>]\n" % (name, label))
    else:
        outfile.write("%s [label = \"%s\"]\n" % (name, label))
    outfile.write("%s [color=%s,style=filled]\n" % (name, color))


done_deps = []


# Print the dependency graph of a package
def print_pkg_deps(outfile, dict_deps, dict_version, stop_list, exclude_list,
                   arrow_dir, draw_graph, depth, max_depth, pkg, colors):
    if pkg in done_deps:
        return
    done_deps.append(pkg)
    if draw_graph:
        print_attrs(outfile, pkg, dict_version.get(pkg), depth, colors)
    elif depth != 0:
        outfile.write("%s " % pkg)
    if pkg not in dict_deps:
        return
    for p in stop_list:
        if fnmatch(pkg, p):
            return
    if dict_version.get(pkg) == "virtual" and "virtual" in stop_list:
        return
    if pkg.startswith("host-") and "host" in stop_list:
        return
    if max_depth == 0 or depth < max_depth:
        for d in dict_deps[pkg]:
            if dict_version.get(d) == "virtual" \
               and "virtual" in exclude_list:
                continue
            if d.startswith("host-") \
               and "host" in exclude_list:
                continue
            add = True
            for p in exclude_list:
                if fnmatch(d, p):
                    add = False
                    break
            if add:
                if draw_graph:
                    outfile.write("%s -> %s [dir=%s]\n" % (pkg_node_name(pkg), pkg_node_name(d), arrow_dir))
                print_pkg_deps(outfile, dict_deps, dict_version, stop_list, exclude_list,
                               arrow_dir, draw_graph, depth + 1, max_depth, d, colors)


def parse_args():
    parser = argparse.ArgumentParser(description="Graph packages dependencies")
    parser.add_argument("--check-only", "-C", dest="check_only", action="store_true", default=False,
                        help="Only do the dependency checks (circular deps...)")
    parser.add_argument("--outfile", "-o", metavar="OUT_FILE", dest="outfile",
                        help="File in which to generate the dot representation")
    parser.add_argument("--package", '-p', metavar="PACKAGE",
                        help="Graph the dependencies of PACKAGE")
    parser.add_argument("--depth", '-d', metavar="DEPTH", dest="depth", type=int, default=0,
                        help="Limit the dependency graph to DEPTH levels; 0 means no limit.")
    parser.add_argument("--stop-on", "-s", metavar="PACKAGE", dest="stop_list", action="append",
                        help="Do not graph past this package (can be given multiple times)." +
                        " Can be a package name or a glob, " +
                        " 'virtual' to stop on virtual packages, or " +
                        "'host' to stop on host packages.")
    parser.add_argument("--exclude", "-x", metavar="PACKAGE", dest="exclude_list", action="append",
                        help="Like --stop-on, but do not add PACKAGE to the graph.")
    parser.add_argument("--colors", "-c", metavar="COLOR_LIST", dest="colors",
                        default="lightblue,grey,gainsboro",
                        help="Comma-separated list of the three colors to use" +
                        " to draw the top-level package, the target" +
                        " packages, and the host packages, in this order." +
                        " Defaults to: 'lightblue,grey,gainsboro'")
    parser.add_argument("--transitive", dest="transitive", action='store_true',
                        default=False)
    parser.add_argument("--no-transitive", dest="transitive", action='store_false',
                        help="Draw (do not draw) transitive dependencies")
    parser.add_argument("--direct", dest="direct", action='store_true', default=True,
                        help="Draw direct dependencies (the default)")
    parser.add_argument("--reverse", dest="direct", action='store_false',
                        help="Draw reverse dependencies")
    parser.add_argument("--quiet", '-q', dest="quiet", action='store_true',
                        help="Quiet")
    parser.add_argument("--flat-list", '-f', dest="flat_list", action='store_true', default=False,
                        help="Do not draw graph, just print a flat list")
    return parser.parse_args()


def main():
    args = parse_args()

    check_only = args.check_only

    logging.basicConfig(stream=sys.stderr, format='%(message)s',
                        level=logging.WARNING if args.quiet else logging.INFO)

    if args.outfile is None:
        outfile = sys.stdout
    else:
        if check_only:
            logging.error("don't specify outfile and check-only at the same time")
            sys.exit(1)
        outfile = open(args.outfile, "w")

    if args.package is None:
        mode = MODE_FULL
    else:
        mode = MODE_PKG
        rootpkg = args.package

    if args.stop_list is None:
        stop_list = []
    else:
        stop_list = args.stop_list

    if args.exclude_list is None:
        exclude_list = []
    else:
        exclude_list = args.exclude_list

    if args.direct:
        get_depends_func = brpkgutil.get_depends
        arrow_dir = "forward"
    else:
        if mode == MODE_FULL:
            logging.error("--reverse needs a package")
            sys.exit(1)
        get_depends_func = brpkgutil.get_rdepends
        arrow_dir = "back"

    draw_graph = not args.flat_list

    # Get the colors: we need exactly three colors,
    # so no need not split more than 4
    # We'll let 'dot' validate the colors...
    colors = args.colors.split(',', 4)
    if len(colors) != 3:
        logging.error("Error: incorrect color list '%s'" % args.colors)
        sys.exit(1)

    # In full mode, start with the result of get_targets() to get the main
    # targets and then use get_all_depends() for all targets
    if mode == MODE_FULL:
        targets = get_targets()
        dependencies = []
        allpkgs.append('all')
        filtered_targets = []
        for tg in targets:
            dependencies.append(('all', tg))
            filtered_targets.append(tg)
        deps = get_all_depends(filtered_targets, get_depends_func)
        if deps is not None:
            dependencies += deps
        rootpkg = 'all'

    # In pkg mode, start directly with get_all_depends() on the requested
    # package
    elif mode == MODE_PKG:
        dependencies = get_all_depends([rootpkg], get_depends_func)

    # Make the dependencies a dictionnary { 'pkg':[dep1, dep2, ...] }
    dict_deps = {}
    for dep in dependencies:
        if dep[0] not in dict_deps:
            dict_deps[dep[0]] = []
        dict_deps[dep[0]].append(dep[1])

    check_circular_deps(dict_deps)
    if check_only:
        sys.exit(0)

    dict_deps = remove_extra_deps(dict_deps, rootpkg, args.transitive)
    dict_version = brpkgutil.get_version([pkg for pkg in allpkgs
                                          if pkg != "all" and not pkg.startswith("root")])

    # Start printing the graph data
    if draw_graph:
        outfile.write("digraph G {\n")

    print_pkg_deps(outfile, dict_deps, dict_version, stop_list, exclude_list,
                   arrow_dir, draw_graph, 0, args.depth, rootpkg, colors)

    if draw_graph:
        outfile.write("}\n")
    else:
        outfile.write("\n")


if __name__ == "__main__":
    sys.exit(main())