#!/usr/bin/env python3

from graphviz import Digraph
import sys
from yaml import load, YAMLError
try:
    from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
    from yaml import Loader, Dumper


def main(args):
    try:
        if len(args) > 0:
            try:
                stream = open(args[0], 'r')
            except IOError:
                print("File not found or path is incorrect")
                return 1
            content = load(stream, Loader=Loader)
            stream.close()
        else:
            # read yaml content from stdout
            content = load(sys.stdin.read(), Loader=Loader)
    except YAMLError:
        print('Invalid yaml file format.')
        return 1

    dot = Digraph('services', comment='services', format='png')
    if 'services' in content:
        for service_group_name in content['services']:
            service_group = content['services'][service_group_name]
            with dot.subgraph(name=f"cluster_{service_group_name}") as sub:
                sub.attr(label=service_group_name)
                if service_group is not None:
                    for service_name in service_group:
                        service = service_group[service_name]
                        if service is not None:  # check for None
                            label = service.get('description', service_name)
                            attrs = service.get('attrs', {})
                            sub.node(service_name, label, attrs)
                            if 'depends_on' in service:
                                dependencies = service['depends_on']
                                if dependencies is not None:
                                    for dependency in dependencies:
                                        # if dependency between services within
                                        # the same cluster
                                        if dependency in service_group:
                                            sub.edge(service_name, dependency)
                                        else:
                                            dot.edge(service_name, dependency)
    dot.render('output/services.gv', view=True)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))