from argparse import ArgumentParser
from pkg_resources import working_set
from graphviz import Digraph


def view_deps_graph(package):
    visited = set()
    visited.add(str(package).split()[0])
    deps_graph = Digraph(f"{str(package).split()[0]}_deps", filename=f"{str(package).split()[0]}_deps.gv")
    dfs(package, visited, deps_graph)


def dfs(current_node, visited: set, deps_graph):
    for dependency in [str(d).split(">=")[0].split("==")[0].split("<=")[0] for d in current_node.requires()]:
        if dependency not in visited:
            deps_graph.edge(dependency, str(current_node).split()[0])
            visited.add(dependency)
            dfs(working_set.by_key[dependency], visited, deps_graph)


def main():
    parser = ArgumentParser(description="Get the graph of dependencies for any python package in GraphViz format")
    parser.add_argument("package_name", help="Name of the package to get dependencies")
    package_name = parser.parse_args().package_name
    package = working_set.by_key[package_name]
    view_deps_graph(package)

if __name__ == "__main__":
    main()