#!/usr/bin/python

import os
import sys
import json
import argparse
from collections import OrderedDict

TEMP_NETCFG_FILE = '/tmp/bmv2-demo-cfg.json'
BASE_LONGITUDE = -115
SWITCH_BASE_LATITUDE = 25
HOST_BASE_LATITUDE = 28
BASE_SHIFT = 8
VLAN_NONE = -1
DEFAULT_SW_BW = 50
DEFAULT_HOST_BW = 25

if 'ONOS_ROOT' not in os.environ:
    print "Environment var $ONOS_ROOT not set"
    exit()
else:
    ONOS_ROOT = os.environ["ONOS_ROOT"]
    sys.path.append(ONOS_ROOT + "/tools/dev/mininet")
if 'RUN_PACK_PATH' not in os.environ:
    print "Environment var $RUN_PACK_PATH not set"
    exit()
else:
    RUN_PACK_PATH = os.environ["RUN_PACK_PATH"]

from onos import ONOSCluster, ONOSCLI
from bmv2 import ONOSBmv2Switch, ONOSHost

from itertools import combinations
from time import sleep
from subprocess import call

from mininet.cli import CLI
from mininet.link import TCLink
from mininet.log import setLogLevel
from mininet.net import Mininet
from mininet.node import RemoteController, Host
from mininet.topo import Topo


class ClosTopo(Topo):
    "2 stage Clos topology"

    def __init__(self, args, **opts):
        # Initialize topology and default options
        Topo.__init__(self, **opts)

        bmv2SwitchIds = []
        for row in (1, 2):
            for col in range(1, args.size + 1):
                bmv2SwitchIds.append("s%d%d" % (row, col))

        bmv2Switches = {}

        for switchId in bmv2SwitchIds:
            deviceId = int(switchId[1:])
            # Use first number in device id to calculate latitude (row number)
            latitude = SWITCH_BASE_LATITUDE + (deviceId // 10) * BASE_SHIFT

            # Use second number in device id to calculate longitude (column number)
            longitude = BASE_LONGITUDE + (deviceId % 10) * BASE_SHIFT
            bmv2Switches[switchId] = self.addSwitch(switchId,
                                                    cls=ONOSBmv2Switch,
                                                    loglevel="warn",
                                                    deviceId=deviceId,
                                                    netcfg=False,
                                                    longitude=longitude,
                                                    latitude=latitude,
                                                    pipeconfId=args.pipeconf_id)

        for i in range(1, args.size + 1):
            for j in range(1, args.size + 1):
                if i == j:
                    # 2 links
                    self.addLink(bmv2Switches["s1%d" % i], bmv2Switches["s2%d" % j],
                                 cls=TCLink, bw=DEFAULT_SW_BW)
                    if args.with_imbalanced_striping:
                        self.addLink(bmv2Switches["s1%d" % i], bmv2Switches["s2%d" % j],
                                     cls=TCLink, bw=DEFAULT_SW_BW)
                else:
                    self.addLink(bmv2Switches["s1%d" % i], bmv2Switches["s2%d" % j],
                                 cls=TCLink, bw=DEFAULT_SW_BW)

        for hostId in range(1, args.size + 1):
            host = self.addHost("h%d" % hostId,
                                cls=DemoHost,
                                ip="10.0.0.%d/24" % hostId,
                                mac='00:00:00:00:00:%02x' % hostId)
            self.addLink(host, bmv2Switches["s1%d" % hostId], cls=TCLink, bw=DEFAULT_HOST_BW)


class DemoHost(ONOSHost):
    "Demo host"

    def __init__(self, name, **params):
        ONOSHost.__init__(self, name, **params)
        self.exectoken = "/tmp/mn-exec-token-host-%s" % name
        self.cmd("touch %s" % self.exectoken)

    def startPingBg(self, h):
        self.cmd(self.getInfiniteCmdBg("ping -i0.5 %s" % h.IP()))
        self.cmd(self.getInfiniteCmdBg("arping -w5000000 %s" % h.IP()))

    def startIperfServer(self):
        self.cmd(self.getInfiniteCmdBg("iperf -s -u"))

    def startIperfClient(self, h, flowBw="512k", numFlows=5, duration=5):
        iperfCmd = "iperf -c{} -u -b{} -P{} -t{}".format(h.IP(), flowBw, numFlows, duration)
        self.cmd(self.getInfiniteCmdBg(iperfCmd, sleep=0))

    def stop(self):
        self.cmd("killall iperf")
        self.cmd("killall ping")
        self.cmd("killall arping")

    def describe(self):
        print "**********"
        print self.name
        print "default interface: %s\t%s\t%s" % (
            self.defaultIntf().name,
            self.defaultIntf().IP(),
            self.defaultIntf().MAC()
        )
        print "**********"

    def getInfiniteCmdBg(self, cmd, logfile="/dev/null", sleep=1):
        return "(while [ -e {} ]; " \
               "do {}; " \
               "sleep {}; " \
               "done;) > {} 2>&1 &".format(self.exectoken, cmd, sleep, logfile)

    def getCmdBg(self, cmd, logfile="/dev/null"):
        return "{} > {} 2>&1 &".format(cmd, logfile)


def generateNetcfg(onosIp, net, args):
    netcfg = OrderedDict()
    netcfg['devices'] = {}
    netcfg['links'] = {}
    netcfg['hosts'] = {}
    # Device configs
    for sw in net.switches:
        srcIp = sw.getSourceIp(onosIp)
        netcfg['devices'][sw.onosDeviceId] = sw.getDeviceConfig(srcIp)

    hostLocations = {}
    # Link configs
    for link in net.links:
        switchPort = link.intf1.name.split('-')
        sw1Name = switchPort[0]  # s11
        port1Name = switchPort[1]  # eth0
        port1 = port1Name[3:]
        switchPort = link.intf2.name.split('-')
        sw2Name = switchPort[0]
        port2Name = switchPort[1]
        port2 = port2Name[3:]
        sw1 = net[sw1Name]
        sw2 = net[sw2Name]
        if isinstance(sw1, Host):
            # record host location and ignore it
            # e.g. {'h1': 'device:bmv2:11'}
            hostLocations[sw1.name] = '%s/%s' % (sw2.onosDeviceId, port2)
            continue

        if isinstance(sw2, Host):
            # record host location and ignore it
            # e.g. {'h1': 'device:bmv2:11'}
            hostLocations[sw2.name] = '%s/%s' % (sw1.onosDeviceId, port1)
            continue

        for linkId in ('%s/%s-%s/%s' % (sw1.onosDeviceId, port1, sw2.onosDeviceId, port2),
                       '%s/%s-%s/%s' % (sw2.onosDeviceId, port2, sw1.onosDeviceId, port1)):
            netcfg['links'][linkId] = {
                'basic': {
                    'type': 'DIRECT',
                    'bandwidth': DEFAULT_SW_BW
                }
            }

    # Host configs
    longitude = BASE_LONGITUDE
    for host in net.hosts:
        longitude = longitude + BASE_SHIFT
        hostDefaultIntf = host.defaultIntf()
        hostMac = host.MAC(hostDefaultIntf)
        hostIp = host.IP(hostDefaultIntf)
        hostId = '%s/%d' % (hostMac, VLAN_NONE)
        location = hostLocations[host.name]

        # use host Id to generate host location
        hostConfig = {
            'basic': {
                'locations': [location],
                'ips': [hostIp],
                'name': host.name,
                'latitude': HOST_BASE_LATITUDE,
                'longitude': longitude
            }
        }
        netcfg['hosts'][hostId] = hostConfig

    netcfg["apps"] = {
        "org.onosproject.core": {
            "core": {
                "linkDiscoveryMode": "STRICT"
            }
        }
    }

    print "Writing network config to %s" % TEMP_NETCFG_FILE
    with open(TEMP_NETCFG_FILE, 'w') as tempFile:
        json.dump(netcfg, tempFile, indent=4)


def main(args):
    if not args.onos_ip:
        controller = ONOSCluster('c0', 3)
        onosIp = controller.nodes()[0].IP()
    else:
        controller = RemoteController('c0', ip=args.onos_ip)
        onosIp = args.onos_ip

    topo = ClosTopo(args)

    net = Mininet(topo=topo, build=False, controller=[controller])

    net.build()
    net.start()

    print "Network started"

    # Generate background traffic.
    sleep(3)
    for (h1, h2) in combinations(net.hosts, 2):
        h1.startPingBg(h2)
        h2.startPingBg(h1)

    print "Background ping started"

    for h in net.hosts:
        h.startIperfServer()

    print "Iperf servers started"

    # sleep(4)
    # print "Starting traffic from h1 to h3..."
    # net.hosts[0].startIperfClient(net.hosts[-1], flowBw="200k", numFlows=100, duration=10)

    generateNetcfg(onosIp, net, args)

    if args.netcfg_sleep > 0:
        print "Waiting %d seconds before pushing config to ONOS..." % args.netcfg_sleep
        sleep(args.netcfg_sleep)

    print "Pushing config to ONOS..."
    call(("%s/onos-netcfg" % RUN_PACK_PATH, onosIp, TEMP_NETCFG_FILE))

    if not args.onos_ip:
        ONOSCLI(net)
    else:
        CLI(net)

    net.stop()
    call(("rm", "-f", TEMP_NETCFG_FILE))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='BMv2 mininet demo script (2-stage Clos topology)')
    parser.add_argument('--onos-ip', help='ONOS-BMv2 controller IP address',
                        type=str, action="store", required=False)
    parser.add_argument('--size', help='Number of leaf/spine switches',
                        type=int, action="store", required=False, default=2)
    parser.add_argument('--with-imbalanced-striping', help='Topology with imbalanced striping',
                        type=bool, action="store", required=False, default=False)
    parser.add_argument('--pipeconf-id', help='Pipeconf ID for switches',
                        type=str, action="store", required=False, default='')
    parser.add_argument('--netcfg-sleep', help='Seconds to wait before pushing config to ONOS',
                        type=int, action="store", required=False, default=5)
    args = parser.parse_args()
    setLogLevel('info')
    main(args)