#!/usr/bin/env python3
"""Prints a PlantUML diagram that shows the DAG of the GitLab pipeline"""

import yaml
from pprint import pprint

def merge(user, default):
    if isinstance(user,dict) and isinstance(default,dict):
        for k,v in default.items():
            if k not in user:
                user[k] = v
            else:
                user[k] = merge(user[k],v)
    return user

merged = {}
with open('.gitlab-ci.yml') as file:
    tmp = yaml.load(file)
    filenames = tmp['include']
    filenames.append('.gitlab-ci.yml')
    
for filename in filenames:
    with open(filename) as file:
        tmp = yaml.load(file)
        merge(merged, tmp)
        tmp = {}

print('@startuml')
print('left to right direction')
jobs = {k: v for k, v in merged.items() if k not in ['image', 'services', 'stages', 'include']}
# This groups by stage, which looks nice but may be misleading
#for stage in merged['stages']:
#    matching_jobs = [k for k, v in jobs.items() if v['stage'] == stage]
#    print(f'partition "{stage}" {{')
#    for job in matching_jobs:
#        needs = merged[job].get('needs')
#        if needs:
#            for need in needs:
#                print(f'  "{need}" --> "{job}" ')
#        else:
#            print(f'  (*) --> "{job}" ')
#    print("}")
for job in jobs:
    needs = merged[job].get('needs')
    if needs:
        for need in needs:
            print(f'"{need}" --> "{job}" ')
    else:
        print(f'(*) --> "{job}" ')
print('@enduml')