#!/usr/bin/env python

from rumba.model import *
from rumba.utils import ExperimentManager

# import testbed plugins
import rumba.testbeds.emulab as emulab
import rumba.testbeds.jfed as jfed
import rumba.testbeds.qemu as qemu

# import prototype plugins
import rumba.prototypes.ouroboros as our
import rumba.prototypes.rlite as rl
import rumba.prototypes.irati as irati

import argparse
import sys


description = "Script to create the demo from the Rumba paper"

argparser = argparse.ArgumentParser(description = description)
argparser.add_argument('-m', '--metro', type = int, default = '2',
                       help = "Number of metro networks")
argparser.add_argument('-n', '--metro-nodes', type = int, default = '5',
                       help = "Number of nodes per metro")
argparser.add_argument('-e', '--end-nodes', type = int, default = '2',
                       help = "Number of nodes connected to each metro")

args = argparser.parse_args()

log.set_logging_level('DEBUG')

pi = NormalDIF("pi")
pi.dt_strategy = 'minimal'

c = NormalDIF("core")

nodes = []

core_nodes = []

for i in range(0, args.metro):
    m = NormalDIF("m-" + str(i))
    e_prev = None
    m1 = None
    mn = None
    for j in range(0, args.metro_nodes):
        if j is not (args.metro_nodes - 1):
            e = ShimEthDIF("e-" + str(i) + "-" + str(j))
        else:
            e = None
        if e_prev == None and e != None:
            node = Node("metro-" + str(i) + "-" + str(j),
                        difs = [m, pi, e],
                        dif_registrations = {pi : [m], m : [e]})
            m1 = node
        elif e_prev != None and e != None:
            node = Node("metro-" + str(i) + "-" + str(j),
                        difs = [m, pi, e, e_prev],
                        dif_registrations = {pi : [m], m : [e, e_prev]})
        else:
            node = Node("metro-" + str(i) + "-" + str(j),
                        difs = [m, pi, e_prev],
                        dif_registrations = {pi : [m], m : [e_prev]})
            mn = node
        nodes.append(node)
        e_prev = e
        # Create nodes attached to metro
        for k in range(0, args.end_nodes):
            ee = ShimEthDIF("e-" + str(i) + "-" + str(j) + "-" + str(k))
            end_node = Node("end-" + str(i) + "-" + str(j) + "-" + str(k),
                            difs = [m, pi, ee],
                            dif_registrations = {pi : [m], m : [ee]})
            node.add_dif(ee)
            node.add_dif_registration(m, ee)
            nodes.append(end_node)

    e = ShimEthDIF("e-" + str(i) + "-" + str(args.metro_nodes))
    m1.add_dif(e)
    m1.add_dif_registration(m, e)
    mn.add_dif(e)
    mn.add_dif_registration(m, e)

    # Add 2 links from metro to core
    e0 = ShimEthDIF("c-e-" + str(i) + "-0")
    e1 = ShimEthDIF("c-e-" + str(i) + "-1")
    e2 = ShimEthDIF("c-e-" + str(i) + "-2")

    c1 = Node("c-" + str(i) + "-0",
              difs = [c, m, pi, e0, e2],
              dif_registrations = {pi : [m, c], m : [e0], c : [e2]})

    c2 = Node("c-" + str(i) + "-1",
              difs = [c, m, pi, e1, e2],
              dif_registrations = {pi : [m, c], m : [e1], c : [e2]})

    nodes.append(c1)
    nodes.append(c2)

    m1.add_dif(e0)
    m1.add_dif_registration(m, e0)
    mn.add_dif(e1)
    mn.add_dif_registration(m, e1)

    for x in range(0, len(core_nodes), 2):
        ce0 = ShimEthDIF("c-c-e-" + str(i) + "-" + str(x))
        ce1 = ShimEthDIF("c-c-e-" + str(i) + "-" + str(x + 1))

        core_nodes[x].add_dif(ce0)
        core_nodes[x].add_dif_registration(c, ce0)
        core_nodes[x + 1].add_dif(ce1)
        core_nodes[x + 1].add_dif_registration(c, ce1)

        c1.add_dif(ce0)
        c1.add_dif_registration(c, ce0)
        c2.add_dif(ce1)
        c2.add_dif_registration(c, ce1)

    core_nodes.append(c1)
    core_nodes.append(c2)

tb = jfed.Testbed(exp_name = "metro2",
                  cert_file = '/home/sander/cert.pem',
                  authority = 'exogeni.net:umassvmsite',
                  username = 'sander')

exp = our.Experiment(tb, nodes = nodes)

with ExperimentManager(exp, do_swap_out = ExperimentManager.PROMPT):
    exp.swap_in()
    exp.install_prototype()
    exp.bootstrap_prototype()