from rumba.model import Node, UnicastLayer, EthDixLayer
import rumba.log as log
# import testbed plugins
import rumba.testbeds.local as local
from rumba.storyboard import *
# import Ouroboros prototype plugin
import rumba.prototypes.ouroboros as our


__all__ = ["local_exp", "name_nodes", "start_net"]


log.set_logging_level("DEBUG")


def buildRack(nRack, numHosts=2):
    rack = UnicastLayer(f"rack{nRack}")
    cables, hosts = [], []

    for n in range(numHosts):
        cables.append(EthDixLayer(f"e{n}r{nRack}"))
        hosts.append(
            Node(
                f"h{n}r{nRack}",
                layers=[cables[-1], rack],
                registrations={rack: [cables[-1]]},
            )
        )

    ToR_switch = Node(
        f"s{nRack}", layers=[*cables, rack], registrations={rack: cables}
    )

    return [*hosts, ToR_switch]  # list of nodes, ToR_switch latest


def build(numRacks=2, numHostsPerRack=2):
    nodes, router_cables = [], []

    for i in range(numRacks):
        router_cables.append(EthDixLayer(f"router_{i}"))
        nodes += buildRack(i, numHosts=numHostsPerRack)
        nodes[-1].add_layer(router_cables[-1])
        #nodes[-1].add_registration(router_cables[-1])

    Datacenter_router = Node("router", layers=router_cables)

    return [*nodes, Datacenter_router]


nodes = build()

name_nodes = list(map(lambda node: node.name, nodes))

local_tb = local.Testbed()


local_exp = our.Experiment(local_tb, nodes=nodes)


def start_net():
    local_exp.swap_in()
    local_exp.install_prototype()
    local_exp.bootstrap_prototype()


def print_network():
    local_exp.export_connectivity_graph("datacenter_physical.pdf")

    for layer in local_exp.layer_ordering:
        if not isinstance(layer, EthDixLayer):
            local_exp.export_layer_graph(f"datacenter_layer_{layer.name}.pdf", layer)


if __name__ == '__main__':
    #start_net()
    print_network()