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

# import testbed plugins
import rumba.testbeds.emulab as emulab
import rumba.testbeds.jfed as jfed
import rumba.testbeds.local as local
import rumba.testbeds.dockertb as docker

# import prototype plugins
import rumba.prototypes.ouroboros as our

import argparse
import json

from rumba.visualizer import draw_network, get_network_from_rumba_experiment

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

argparser = argparse.ArgumentParser(description=description)

# Parameters to set the experiment size
argparser.add_argument('--metro-networks', type=int, default='2',
                       help="Number of metro networks")
argparser.add_argument('--metro-nodes', type=int, default='5',
                       help="Number of nodes per metro")
argparser.add_argument('--end-nodes', type=int, default='2',
                       help="Number of nodes connected to each metro")

# Other parameters
argparser.add_argument('--verbosity', metavar='VERBOSITY', type=str,
                       default='INFO',
                       choices=['DEBUG', 'INFO', 'WARNING',
                                'ERROR', 'CRITICAL'],
                       help='Logging verbosity level')
argparser.add_argument('--swapout', type=SwapOutStrategy.from_string,
                       default=PROMPT_SWAPOUT,
                       choices=list(SwapOutStrategy),
                       help='What action to perform on swap-out')
argparser.add_argument('--visualize', action="store_true",
                       help='Draw a representation of the network')


subparsers = argparser.add_subparsers(dest='testbed')
emulab_p = subparsers.add_parser('emulab', help='Use emulab testbed')
jfed_p = subparsers.add_parser('jfed', help='Use jfed testbed')
local_p = subparsers.add_parser('local', help='Use local testbed')
docker_p = subparsers.add_parser('docker', help='Use Docker testbed')

emulab_p.add_argument('--url', metavar='URL', type=str,
                      default="wall2.ilabt.iminds.be",
                      help='Url')
emulab_p.add_argument('--image', metavar='IMG', type=str,
                      default="UBUNTU14-64-STD",
                      help='Ubuntu image')
emulab_p.add_argument('--username', type=str, required=True,
                      help="Username")
emulab_p.add_argument('--password', type=str, help="Password")
emulab_p.add_argument('--exp-name', type=str, required=True,
                      help="Name of the experiment")

jfed_p.add_argument('--cert-file', type=str,
                    required=True,
                    help="Absolute path to certificate (.pem) file"
                    " to be used with jFed")
jfed_p.add_argument('--authority', type=str, default='exogeni.net',
                    help="jFed testbed to use")
jfed_p.add_argument('--image', metavar='IMAGE', type=str,
                    default=None,
                    help='Image to be used')
jfed_p.add_argument('--exp-hours', metavar='HOURS', type=str,
                    default="2", help='Experiment hours')
jfed_p.add_argument('--username', type=str, required=True,
                    help="Username")
jfed_p.add_argument('--exp-name', type=str, required=True,
                    help="Name of the experiment")

docker_p.add_argument('--base-image', metavar='DOCKIMAGE', type=str,
                      default=None, required=True,
                      help='Docker image to be used')
docker_p.add_argument('--use-ovs', action='store_true',
                      default=False, help='Use ovs switch')

args = argparser.parse_args()

log.set_logging_level(args.verbosity)

pi = UnicastLayer("pi")

c = UnicastLayer("core")

nodes = []

core_nodes = []

for i in range(0, args.metro_networks):
    m = UnicastLayer("m-" + str(i))
    e_prev = None
    m1 = None
    mn = None
    for j in range(0, args.metro_nodes):
        if j != (args.metro_nodes - 1):
            e = EthDixLayer("e-" + str(i) + "-" + str(j))
        else:
            e = None
        if e_prev is None and e is not None:
            node = Node("metro-" + str(i) + "-" + str(j),
                        layers=[m, pi, e],
                        registrations={pi: [m], m: [e]})
            m1 = node
        elif e_prev is not None and e is not None:
            node = Node("metro-" + str(i) + "-" + str(j),
                        layers=[m, pi, e, e_prev],
                        registrations={pi: [m], m: [e, e_prev]})
        else:
            node = Node("metro-" + str(i) + "-" + str(j),
                        layers=[m, pi, e_prev],
                        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 = EthDixLayer("e-" + str(i) + "-" + str(j) + "-" + str(k))
            end_node = Node("end-" + str(i) + "-" + str(j) + "-" + str(k),
                            layers=[pi, ee],
                            registrations={pi: [ee]})
            node.add_layer(ee)
            node.add_registration(pi, ee)
            nodes.append(end_node)

    e = EthDixLayer("e-" + str(i) + "-" + str(args.metro_nodes))
    m1.add_layer(e)
    m1.add_registration(m, e)
    mn.add_layer(e)
    mn.add_registration(m, e)

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

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

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

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

    m1.add_layer(e0)
    m1.add_registration(m, e0)
    mn.add_layer(e1)
    mn.add_registration(m, e1)

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

        core_nodes[x].add_layer(ce0)
        core_nodes[x].add_registration(c, ce0)
        core_nodes[x + 1].add_layer(ce1)
        core_nodes[x + 1].add_registration(c, ce1)

        c1.add_layer(ce0)
        c1.add_registration(c, ce0)
        c2.add_layer(ce1)
        c2.add_registration(c, ce1)

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

if args.testbed == 'emulab':
    test_class = emulab.Testbed
    testbed_args = {a.dest: getattr(args, a.dest)
                    for a in emulab_p._actions
                    if a.dest != 'help'
                    and getattr(args, a.dest) is not None}
elif args.testbed == 'jfed':
    test_class = jfed.Testbed
    testbed_args = {a.dest: getattr(args, a.dest)
                    for a in jfed_p._actions
                    if a.dest != 'help'
                    and getattr(args, a.dest) is not None}
elif args.testbed == 'local':
    test_class = local.Testbed
    testbed_args = {a.dest: getattr(args, a.dest)
                    for a in local_p._actions
                    if a.dest != 'help'
                    and getattr(args, a.dest) is not None}
elif args.testbed == 'docker':
    test_class = docker.Testbed
    testbed_args = {a.dest: getattr(args, a.dest)
                    for a in docker_p._actions
                    if a.dest != 'help'
                    and getattr(args, a.dest) is not None}
elif args.testbed is None:
    print('Testbed type must be specified!')
    print(argparser.format_help())
    exit(1)

testbed = test_class(**testbed_args)

exp_class = our.Experiment

exp = exp_class(testbed, nodes=nodes)

if __name__ == "__main__":

    if args.visualize:
        nw = get_network_from_rumba_experiment(exp)
        draw_network(nw)
    with ExperimentManager(exp, swap_out_strategy=args.swapout) as expM:
        exp.swap_in()
        if not isinstance(testbed, docker.Testbed) \
                and not isinstance(testbed, local.Testbed):
            exp.install_prototype()
        exp.bootstrap_prototype()

    exit(expM.exit_status)
