import os

from pyvis.network import Network

from steps import get_step_depends_on, steps


LEVELS = (
    "artist.prepare",
    "track.prepare",
    "artist.genre",
    "track.genre",
    ".extract",
    ".transform",
    ".load",
    ".switch_ready",
    ".switch",
    ".cleanup",
    "validation.",
    "upload_cms_data."
)
FILE_NAME = "graph.html"


def add_node(net_, name_):
    parts = name_.split(".")
    group = ".".join(parts[:2])
    level = 0
    for level_ in LEVELS:
        if level_ in name_:
            level = LEVELS.index(level_)
            break
    net_.add_node(name_, group=group, level=level, shape="box")


if __name__ == "__main__":
    net = Network(directed=True, height="1000px", layout=True)
    setattr(net.options.layout.hierarchical, "nodeSpacing", 280)
    net.toggle_physics(False)
    for step_cls in steps:
        add_node(net, step_cls.step_name)
        for parent in get_step_depends_on(step_cls):
            add_node(net, parent.step_name)
            add_node(net, step_cls.step_name)
            net.add_edge(parent.step_name, step_cls.step_name)

    net.write_html(FILE_NAME)
    print(f"Graph generated {os.path.join(os.getcwd(), FILE_NAME)}")