import sys

from tests.local.graph_connect import graph_client

g = graph_client.graph_traversal

parallelism = int(sys.argv[1])

components = {"Sp" : parallelism, "A" : parallelism, "B" : parallelism,
              "C" : parallelism}

task_id = 0

# Create all the nodes
print("Creating nodes")

for component, parallelism in components.items():

    print(f"Adding nodes for component {component}")

    for _ in range(parallelism):

        print(f"Creating node: {task_id}")

        (g.addV("instance").property("component", component)
                           .property("task_id", task_id)
                           .next())

        task_id += 1

# Create all the connections
print("Creating edges")
order = list(components.keys())
for i, source_component in enumerate(order):

    if i == (len(order) - 1):
        break

    print(f"Connecting instances of component {source_component}")

    destination_component = order[i+1]

    sources = (g.V().hasLabel("instance")
                    .has("component", source_component)
                    .toList())

    destinations = (g.V().hasLabel("instance")
                         .has("component", destination_component)
                         .toList())

    for j, source in enumerate(sources):
        print(f"Connecting source vertex {j+1} of {len(sources)}")
        for destination in destinations:
            g.V(source).addE("logical").to(destination).next()

print("Done")