#
# A library to manage ARCFIRE experiments
#
#    Copyright (C) 2017-2018 Nextworks S.r.l.
#    Copyright (C) 2017-2018 imec
#
#    Sander Vrijders   <sander.vrijders@ugent.be>
#    Dimitri Staessens <dimitri.staessens@ugent.be>
#    Vincenzo Maffione <v.maffione@nextworks.it>
#    Marco Capitani    <m.capitani@nextworks.it>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., http://www.fsf.org/about/contact/.
#

import time
import subprocess
import re

import rumba.ssh_support as ssh
import rumba.model as mod
import rumba.multiprocess as m_processing
import rumba.log as log
import rumba.testbeds.local as local
import rumba.testbeds.localnet as localnet
import rumba.storyboard as sb
from rumba.elements.topology import LayerType
from rumba.irm_backend import IrmPython, IrmCLI
from rumba.process import Process

logger = log.get_logger(__name__)


class OurServer(sb.Server):

    def __init__(self, server):
        super(OurServer, self).__init__(
            server.ap,
            server.arrival_rate,
            server.actual_parameter + server.min_duration,
            server.options,
            server.max_clients,
            server.clients,
            server.nodes,
            server.min_duration,
            server.id,
            server.as_root,
            server.layers
        )

    def _make_run_cmd(self, node):
        o_cmd = super(OurServer, self)._make_run_cmd(node)

        # Run and store PID
        n_cmd = 'pid=$(%s) && ' % (o_cmd,)

        # Build register command
        r_cmd = 'irm n r %s ' % (self.id,)
        if len(self.layers) == 0:
            r_cmd += ' '.join('ipcp %s' % (ipcp.name,) for ipcp in node.ipcps)
        else:
            for layer in self.layers:
                for ipcp in node.ipcps:
                    if ipcp.layer is layer:
                        r_cmd += 'ipcp %s' % (ipcp.name,)
        r_cmd += ' && '

        # Add register command
        n_cmd += r_cmd

        # Add bind command
        n_cmd += 'irm b process $pid name %s && ' % (self.id,)

        n_cmd += 'echo $pid'  # We need to return the pid for the sb

        return n_cmd


class Experiment(mod.Experiment):
    """
    Represents an Ouroboros experiment.
    """
    def __init__(self,
                 testbed,
                 nodes=None,
                 git_repo='https://ouroboros.rocks/git/ouroboros',
                 git_branch='master',
                 build_options='',
                 add_packages=None,
                 enrollment_strategy='minimal',
                 flows_strategy='full-mesh',
                 influxdb=None):
        """
        Initializes the experiment class.

        :param testbed: The testbed to run the experiment on.
        :param nodes: The list of nodes.
        :param git_repo: The git repository to use for installation.
        :param git_branch: The branch of the git repository to use.
        :param build_options: Options for CMake building
        :param add_packages: Additional packages to be installed.
        :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
        :param flows_strategy: For flows, 'full-mesh', 'minimal' or 'manual'.
        :param influxdb: dict with information for InfluxDB database
            {
                "ip": <IP address for InfluxDB database>,
                "port": <TCP port for InfluxDB database>,
                "org": <Organization in InfluxDB database>,
                "token": <Token to access InfluxDB database>
            }
        """
        mod.Experiment.__init__(
            self,
            testbed,
            nodes,
            git_repo,
            git_branch,
            build_options,
            add_packages,
            enrollment_strategy=enrollment_strategy,
            flows_strategy=flows_strategy,
            server_decorator=OurServer
        )

        if influxdb is not None:
            _required = ['ip', 'port', 'org', 'token']
            for field in _required:
                if field not in influxdb:
                    raise (Exception("Missing field %s in influx_db: %s" % field))

        self.r_ipcps = dict()
        self.irmd = None
        self.influxdb = influxdb
        self.prototype_location = '/tmp/prototype/'
        self.set_startup_command("irmd")
        self.metrics_python_version = "python3.9"

        # Create IRM backend
        if isinstance(testbed, (local.Testbed,
                                localnet.LocalNetworkTestbed)):
            try:
                self.irm = IrmPython()
                logger.info("Using pyouroboros IRM backend")
            except ImportError:
                self.irm = IrmCLI()
                logger.info("pyouroboros not available, "
                            "using CLI IRM backend for local testbed")
        else:
            self.irm = IrmCLI()
            logger.info("Using CLI IRM backend for remote testbed")

    @staticmethod
    def make_executor(node, packages, testbed):
        def executor(commands):
            if packages and len(packages) > 0:
                ssh.aptitude_install(testbed, node, packages)
            if commands and len(commands) > 0:
                node.execute_commands(commands, time_out=None, use_proxy=True)
        return executor

    def prototype_name(self):
        return 'ouroboros'

    def exec_local_cmd(self, cmd):
        try:
            logger.info(cmd)
            subprocess.check_call(cmd.split(' '))
        except subprocess.CalledProcessError as e:
            logger.error("Return code was " + str(e.returncode))
            raise

    def exec_local_cmds(self, cmds):
        for cmd in cmds:
            self.exec_local_cmd(cmd)

    def setup_ouroboros(self):
        try:
            import rumba.testbeds.dockertb as dockertb
            if isinstance(self.testbed, dockertb.Testbed):
                return
        except ImportError:
            pass

        if self._is_local:
            subprocess.check_call('sudo -v'.split())
            popen = subprocess.Popen(["sudo", "irmd"])
            self.irmd = Process(None, popen.pid, "irmd", popen=popen)
            logger.info("Started IRMd (pid %d), sleeping 2 seconds...",
                        popen.pid)
            time.sleep(2)
        else:
            for node in self.nodes:
                node.execute_command("sudo nohup irmd > /dev/null &",
                                     time_out=None)
            time.sleep(1)

    def _install_packages_and_execute_cmds(self, packages, cmds, nodes=None):
        names = []
        executors = []
        args = []

        if nodes and isinstance(nodes, str):
            nodes = [nodes]

        if nodes is None:
            _nodes = self.nodes
        else:
            _nodes = [node for node in self.nodes if node.name in nodes]

        for node in _nodes:
            executor = self.make_executor(node, packages, self.testbed)
            names.append(node.name)
            executors.append(executor)
            args.append(cmds)

        m_processing.call_in_parallel(names, args, executors)

    @property
    def _is_local(self):
        """True if the testbed runs locally (single IRMd)."""
        return isinstance(self.testbed,
                          (local.Testbed, localnet.LocalNetworkTestbed))

    def install_ouroboros(self):
        if self._is_local:
            return

        packages = ["cmake", "protobuf-c-compiler", "git", "libfuse-dev",
                    "libgcrypt20-dev", "libssl-dev"] + self.add_packages

        cmds = ["sudo DEBIAN_FRONTEND=noninteractive apt-get install libprotobuf-c-dev --yes || true",
                "sudo rm -r " + self.prototype_location + " || true",
                "git clone -b " + self.git_branch + " " + self.git_repo + " " + self.prototype_location,
                "cd " + self.prototype_location + " && mkdir build && cd build && " +
                "cmake " + self.build_options + " .. && " +
                "sudo make install -j$(nproc) && " +
                "sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf"]

        self._install_packages_and_execute_cmds(packages, cmds)

    def update_ouroboros(self, branch):
        if self._is_local:
            return

        if branch is not None:
            self.git_branch = branch

        cmds = ["cd " + self.prototype_location + " && " +
                "git checkout " + self.git_branch + " && " +
                "git pull && cd build && sudo make install -j$(nproc)"]

        self._install_packages_and_execute_cmds([], cmds)

    def _install_python_ppa(self):
        packages = ["software-properties-common"]
        python_ppa = "ppa:deadsnakes/ppa"
        cmds = ["sudo add-apt-repository {} -y".format(python_ppa),
                "sudo apt-get update"]

        self._install_packages_and_execute_cmds(packages, cmds)

    def _install_python_version(self, python_version):
        python_packages = ["distutils", "venv"]
        python_libs = ["dev"]

        packages = [python_version]
        for package in python_packages:
            packages += [python_version + "-" + package]

        for lib in python_libs:
            packages += ["lib" + python_version + "-" + lib]

        self._install_packages_and_execute_cmds(packages, [])

    def install_ouroboros_python_exporter(self):
        if self._is_local:
            return

        if self.influxdb is None:
            return

        python_version = self.metrics_python_version
        python_cmd = "venv/bin/" + python_version

        self._install_python_ppa()
        self._install_python_version(python_version)

        cmds = [python_version + " -m venv venv",
                python_cmd + " -m easy_install pip",
                python_cmd + " -m pip install 'influxdb-client[ciso]'"]

        # Clone the metrics exporter repo and get the exporter
        cmds += ["git clone https://ouroboros.rocks/git/ouroboros-metrics",
                 "cp ouroboros-metrics/exporters-influxdb/pyExporter/oexport.py .",
                 "rm -rf ouroboros-metrics"]

        # Write the configuration file
        cmds += ["echo [influx2] > config.ini",
                 "echo 'url=http://" + self.influxdb['ip'] + ":" + str(self.influxdb['port']) + "' >> config.ini",
                 "echo 'org=" + self.influxdb['org'] + "' >> config.ini",
                 "echo 'token=" + self.influxdb['token'] + "' >> config.ini",
                 "echo 'timeout=6000' >> config.ini",
                 'echo "proxy=$http_proxy" >> config.ini',
                 "echo 'verify_ssl=False' >> config.ini"]

        self._install_packages_and_execute_cmds([], cmds)

    def start_ouroboros_metrics_exporter(self, nodes, interval=100):
        cmds = ["nohup ~/venv/bin/python3.9 ~/oexport.py --interval {} > /dev/null 2>&1 &".format(interval)]

        self._install_packages_and_execute_cmds([], cmds, nodes)

    def stop_ouroboros_metrics_exporter(self, nodes):
        cmds = ["kill $(ps ax | grep oexport.py | grep -v grep| cut -f1 -d' ') || true"]

        self._install_packages_and_execute_cmds([], cmds, nodes)

    def _set_phy_link_rate_to(self, src, dst, megabit):

        autoneg = 'off'

        if megabit == 1000:
            autoneg = 'on'

        commands = ["route | grep $(ping server -c 1 | head -n1 | cut -f3 -d' ' | head -c -3 | tail -c +2) | "
                    "grep -o '[^ ]*$' > iface ".format(dst),
                    "sudo ethtool -s $(cat iface) speed {} duplex full autoneg {}".format(megabit, autoneg),
                    "sudo ip link set $(cat iface) down",
                    "sudo ip link set $(cat iface) up && sleep 10",
                    "sudo ethtool $(cat iface) | grep '{}Mb/s'".format(megabit),
                    "rm iface",
                    "while ! ping -c 1 {}; do sleep 1; done".format(dst)]

        self._install_packages_and_execute_cmds([], commands, [src])

    def _get_n1_ipcp_names(self, node, ipcp):
        """Get N-1 IPCP names for name registration."""
        reg_ipcps = []
        if ipcp.layer in node.registrations:
            for lower in node.registrations[ipcp.layer]:
                for ipcp_b in node.ipcps:
                    if ipcp_b in lower.ipcps:
                        reg_ipcps.append(ipcp_b.name)
        return reg_ipcps

    def create_ipcps(self):
        for node in self.nodes:
            for ipcp in node.ipcps:
                layer = ipcp.layer
                layer_type = layer.layer_type

                # Plain local testbed: use local IPCP type for eth layers
                # (LocalNetworkTestbed has real veth interfaces)
                if layer.is_eth \
                   and isinstance(self.testbed, local.Testbed) \
                   and not isinstance(self.testbed,
                                      localnet.LocalNetworkTestbed):
                    layer_type = LayerType.LOCAL

                if layer.is_shim:
                    # Shim IPCPs are always bootstrapped
                    eth_dev = None
                    ip_addr = None
                    if layer.is_eth \
                       and (not isinstance(self.testbed, local.Testbed)
                            or isinstance(self.testbed,
                                          localnet.LocalNetworkTestbed)):
                        eth_dev = ipcp.ifname
                    if layer.is_udp and hasattr(ipcp, 'ip_addr'):
                        ip_addr = ipcp.ip_addr

                    self.irm.create_ipcp(node, ipcp.name, layer_type)
                    self.irm.bootstrap_ipcp(
                        node, ipcp.name, layer_type,
                        layer.name, eth_dev=eth_dev, ip_addr=ip_addr,
                        autobind=False)
                elif ipcp.bootstrapper:
                    # Unicast/broadcast bootstrapper
                    self.irm.create_ipcp(node, ipcp.name, layer_type)

                    policies = {}
                    if hasattr(layer, 'policy'):
                        pols = layer.policy.get_policies()
                        for comp, pol_dict in pols.items():
                            if pol_dict:
                                policies[comp] = next(iter(pol_dict))

                    self.irm.bootstrap_ipcp(
                        node, ipcp.name, layer_type,
                        layer.name, policies=policies)

                    # Postpone N-1 registrations
                    self.r_ipcps[ipcp] = \
                        self._get_n1_ipcp_names(node, ipcp)
                else:
                    # Non-bootstrapper: just create, will enroll later
                    self.irm.create_ipcp(node, ipcp.name, layer_type)

    def enroll_layer(self, el):
        for e in el:
            ipcp = e['enrollee']
            node = ipcp.node

            # Execute postponed registration for enroller
            if e['enroller'] in self.r_ipcps:
                reg_ipcps = self.r_ipcps[e['enroller']]
                enroller_node = e['enroller'].node
                self.irm.reg_name(enroller_node, e['enroller'].name,
                                  reg_ipcps)
                self.irm.reg_name(enroller_node, e['enroller'].layer.name,
                                  reg_ipcps)
                del self.r_ipcps[e['enroller']]

            # Register enrollee name in N-1 layers
            reg_ipcps = self._get_n1_ipcp_names(node, ipcp)
            self.irm.reg_name(node, ipcp.name, reg_ipcps)

            # Enroll
            self.irm.enroll_ipcp(node, ipcp.name, e['enroller'].name)

            # Register layer name
            self.irm.reg_name(node, ipcp.layer.name, reg_ipcps)

    def setup_flows(self, el):
        for e in el:
            ipcp = e['src']
            retry = 0
            max_retries = 3
            while retry < max_retries:
                time.sleep(retry * 5)
                try:
                    self.irm.connect_ipcp(
                        ipcp.node, ipcp.name, e['dst'].name)
                    break
                except Exception as ex:
                    retry += 1
                    logger.error('Failed to connect IPCP, retrying: ' +
                                 str(retry) + '/' + str(max_retries) +
                                 ' retries')
            if retry == max_retries:
                raise Exception('Failed to connect IPCP')

    def _install_prototype(self):
        logger.info("Installing Ouroboros...")
        self.install_ouroboros()
        if self.influxdb:
            logger.info("Installing Ouroboros metrics exporter")
            self.install_ouroboros_python_exporter()
        logger.info("Installed on all nodes...")

    def _bootstrap_prototype(self):
        for layer in self.layer_ordering:
            if hasattr(layer, 'qos_cubes') and len(layer.qos_cubes) != 0:
                logger.warning('QoS cubes not (yet) supported by '
                               'the Ouroboros plugin. Will ignore.')
        logger.info("Starting IRMd on all nodes...")
        self.setup_ouroboros()
        logger.info("Creating IPCPs")
        self.create_ipcps()
        logger.info("Enrolling IPCPs...")

        for enrolls, flows in zip(self.enrollments,
                                  self.flows):
            self.enroll_layer(enrolls)
            self.setup_flows(flows)

        logger.info("All done, have fun!")

    def _update_prototype(self, branch):
        logger.info("Updating Ouroboros")
        self.update_ouroboros(branch)
        logger.info("Update done")

    def _terminate_prototype(self, force=False):
        cmds = list()

        if force is True:
            kill = 'killall -9 '
            cmds.append(kill + 'irmd || true')
            cmds.append(kill + 'ipcpd-unicast || true')
            cmds.append(kill + 'ipcpd-broadcast || true')
            cmds.append(kill + 'ipcpd-eth-llc || true')
            cmds.append(kill + 'ipcpd-eth-dix || true')
            cmds.append(kill + 'ipcpd-udp4 || true')
            cmds.append(kill + 'ipcpd-udp6 || true')
            cmds.append(kill + 'ipcpd-local || true')
            cmds.append(kill + 'ocbr || true')
            cmds.append(kill + 'oping || true')
            cmds.append(kill + 'operf || true')
            cmds.append(kill + 'oecho || true')
            cmds.append(kill + 'ioquake3.x86_64 || true')
            cmds.append(kill + 'ioq3ded.x86_64 || true')
            cmds.append(kill + 'ouroborosio || true')
            cmds.append('rm -rf /dev/shm/ouroboros.* || true')
            cmds.append('bash -c "for i in /tmp/ouroboros/*; do fusermount -u $i; rmdir $i; done || true"')
            cmds.append('rm -rf /tmp/ouroboros || true')
            cmds.append('kill -9 $(ps axjf | grep \'sudo irmd\' | grep -v grep | cut -f4 -d " ") || true')
        else:
            cmds.append('killall -15 irmd')

        logger.info("Killing Ouroboros...")
        if self._is_local:
            cmds = list(map(lambda c: "sudo %s" % (c,), cmds))
            for cmd in cmds:
                subprocess.call(cmd.split())
        else:
            for node in self.nodes:
                node.execute_commands(cmds, time_out=None, as_root=True)

    def _start_metrics_exporter(self, nodes, interval):
        self.start_ouroboros_metrics_exporter(nodes, interval)

    def _stop_metrics_exporter(self, nodes):
        self.stop_ouroboros_metrics_exporter(nodes)

    def destroy_layer(self, layer):
        for ipcp in layer.ipcps:
            self.irm.destroy_ipcp(ipcp.node, ipcp.name)

    def parse_stats(self, lines, spaces=0):
        d = {}

        while len(lines):
            line = lines[0]

            if not re.match(" {%i}.*" % spaces, line):
                return d

            lines.pop(0)

            line = line.strip()

            if re.match(".*:.*", line):
                head, tail = line.split(":", 1)

                if len(tail) == 0:
                    d[head] = self.parse_stats(lines, spaces+1)
                else:
                    d[head] = tail.strip()

        return d

    def export_layer_bandwidth(self, filename, layer):
        f = open(filename, 'w')

        for node in layer.members:
            ipcp = node.get_ipcp_by_layer(layer)

            # Get IPCP address
            if not hasattr(ipcp, 'address'):
                path = '/tmp/ouroboros/' + ipcp.name + '/dt*'
                dt_path = node.execute_command('ls -d %s' % path).stdout
                dts = dt_path.split('.')
                ipcp.address = int(dts[-1])
                logger.info('IPCP %s has dt component '
                            'with address %d' % (ipcp.name, ipcp.address))

        for node in layer.members:
            ipcp = node.get_ipcp_by_layer(layer)

            dt_path = '/tmp/ouroboros/' + ipcp.name + '/dt.' + \
                str(ipcp.address) + '/'

            # Get flows to other endpoints
            fd = node.execute_command('ls --ignore=[01] %s' % dt_path).stdout
            fds = fd.split('\n')
            for fd in fds:
                fd_path = dt_path + fd
                fd_file = node.execute_command('cat %s' % fd_path).stdout

                d = self.parse_stats(fd_file.splitlines())
                remote = d["Endpoint address"]
                ipcp2_name = ''
                for ipcp2 in layer.ipcps:
                    if ipcp2.address == int(remote):
                        ipcp2_name = ipcp2.name

                nr = d["Qos cube   0"]["sent (bytes)"]

                f.write('%s;%s;%s\n' % (ipcp.name, ipcp2_name, nr))

        f.close()
        logger.info('Wrote stats to %s', filename)
