#
# 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 abc
import os
import shutil
import time

import rumba.log as log
import rumba.elements.topology as topology

logger = log.get_logger(__name__)


tmp_dir = '/tmp/rumba'


def _enrollment_root(adj):
    """Pick a graph-central node as the root of the enrollment spanning
    tree, so the tree is as shallow as possible.

    Rooting a line at an endpoint forces the upper layer's directory and
    routing to converge across the maximum number of hops (e.g. a-b-c
    rooted at a leaves c two hops from the bootstrapper); a central root
    halves that.  *adj* maps node -> [(neighbour, lower_layer), ...].
    The node with minimum eccentricity wins; reachability is maximised
    first to stay well-defined on a disconnected graph, and ties fall
    back to insertion (node) order.
    """
    best = None
    best_key = None

    for src in adj:
        dist = {src: 0}
        queue = [src]
        head = 0

        while head < len(queue):
            cur = queue[head]
            head += 1
            for nxt, _lower in adj[cur]:
                if nxt not in dist:
                    dist[nxt] = dist[cur] + 1
                    queue.append(nxt)

        key = (-len(dist), max(dist.values()))
        if best_key is None or key < best_key:
            best_key = key
            best = src

    return best


class Testbed(object):
    """
    Base class for every testbed plugin.
    """

    # Whether this testbed provides interfaces that can carry tc link
    # quality (rate/delay/loss).  Process-only testbeds (e.g. the local
    # testbed) have no interfaces, so applying a LinkQuality is a no-op.
    supports_link_quality = True

    # Whether rumba must build/install the prototype on the nodes.  Testbeds
    # that run on the host against already-installed binaries (local, localnet,
    # nstb) set this False; remote testbeds (emulab, jfed, ...) leave it True.
    needs_install = True

    # Whether the prototype runs as a single instance in one namespace/host:
    # one IRMd, driven as local subprocesses / the in-process IRM backend.
    # Testbeds that give each node its own namespace (nstb) or its own
    # container/host (docker, remote) set this False.
    single_irmd = False

    # Whether swap_out destroys the execution context (e.g. network + mount
    # namespaces) the prototype runs in, so the prototype must be terminated
    # *before* swap_out while it is still reachable.  nstb sets this True.
    terminate_before_swapout = False

    def __init__(self,
                 exp_name,
                 username,
                 password,
                 proj_name,
                 system_logs=None):
        """
        :param exp_name: The experiment name.
        :param username: The username.
        :param password: The password.
        :param proj_name: The project name.
        :param system_logs: Location of the system logs of
                            images of the testbed.
        """
        self.username = username
        self.password = password
        self.proj_name = proj_name
        self.exp_name = exp_name
        self.flags = {'no_vlan_offload': False}
        self.executor = None
        if system_logs is None:
            self.system_logs = ['/var/log/syslog']
        elif isinstance(system_logs, str):
            self.system_logs = [system_logs]
        else:
            self.system_logs = system_logs

    def swap_in(self, experiment):
        """
        Swaps experiment in on the testbed.

        :param experiment: The experiment.
        """
        for node in experiment.nodes:
            node.executor = self.executor

        self._swap_in(experiment)

        # Process-only testbeds have no interfaces to shape; applying a
        # LinkQuality there is a no-op (see Testbed.supports_link_quality).
        if self.supports_link_quality:
            for layer in experiment.layer_ordering:
                if layer.is_eth:
                    layer.link_quality.apply(layer)

    @abc.abstractmethod
    def _swap_in(self, experiment):
        logger.info("_swap_in(): nothing to do")

    def swap_out(self, experiment):
        """
        Swaps experiment out of the testbed.

        :param experiment: The experiment.
        """
        self._swap_out(experiment)

    @abc.abstractmethod
    def _swap_out(self, experiment):
        logger.info("swap_out(): nothing to do")


class Experiment(object):
    """
    Base class for experiments.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, testbed,
                 nodes=None,
                 git_repo=None,
                 git_branch=None,
                 build_options=None,
                 add_packages=None,
                 log_dir=None,
                 prototype_logs=None,
                 enrollment_strategy='minimal',
                 flows_strategy='full-mesh',
                 server_decorator=None):
        """
        :param testbed: The testbed of the experiment.
        :param nodes: The list of nodes in the experiment.
        :param git_repo: The git repository of the prototype.
        :param git_branch: The git branch of the repository.
        :param build_options: Options for building the prototype.
        :param add_packages: Additional packages to be installed.
        :param log_dir: Where to log output of the experiment.
        :param prototype_logs: Where the prototype logs its output.
        :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
        :param flows_strategy: For data flows, 'full-mesh', 'minimal' or 'manual'.
        :param server_decorator: a decorator function which will be applied to
                                 storyboard.server instances when using
                                 this prototype
        """
        if nodes is None:
            nodes = list()
        if add_packages is None:
            add_packages = list()

        self.nodes = nodes
        if server_decorator is None:
            def server_decorator(server):
                return server
        self.server_decorator = server_decorator
        self.git_repo = git_repo
        self.git_branch = git_branch
        self.build_options = build_options
        self.add_packages = add_packages
        self.testbed = testbed
        # the strategy employed for completing the enrollment phase in
        # the different Layers
        self.enrollment_strategy = enrollment_strategy  # 'full-mesh', 'manual'
        # the strategy employed for setting up the data transfer
        # network in the Layers after enrollment
        self.flows_strategy = flows_strategy
        self.layer_ordering = []
        self.enrollments = []  # a list of per-Layer lists of enrollments
        self.flows = []  # a list of per-Layer lists of flows

        if self.enrollment_strategy not in ['full-mesh', 'minimal', 'manual']:
            raise Exception('Unknown enrollment strategy "%s"'
                            % self.enrollment_strategy)
        if self.flows_strategy not in ['full-mesh', 'minimal', 'manual']:
            raise Exception('Unknown flows strategy "%s"'
                            % self.flows_strategy)

        # Determine log directory
        if log_dir is None:
            # If it is None, use /tmp/rumba/{project}
            # Wipe it and make it again
            exp_name = self.testbed.exp_name.replace('/', '_')  # Just in case
            log_dir = os.path.join(tmp_dir, exp_name)
            shutil.rmtree(log_dir, ignore_errors=True)
            os.mkdir(log_dir)
        self.log_dir = log_dir
        if not os.path.isdir(self.log_dir):
            raise Exception('Destination "%s" is not a directory. '
                            'Cannot fetch logs.'
                            % self.log_dir)
        self.prototype_logs = prototype_logs \
            if prototype_logs is not None else []

        # Generate missing information
        self.generate()

    def __repr__(self):
        s = ""
        for n in self.nodes:
            s += "\n" + str(n)

        return s

    def add_node(self, node):
        """
        Adds a node to the experiment.

        :param node: A node.
        """
        self.nodes.append(node)
        self.generate()

    def del_node(self, node):
        """
        Deletes a node from the experiment.

        :param node: A node.
        """
        self.nodes.remove(node)
        self.generate()

    # Compute registration/enrollment order for Layers
    def compute_layer_ordering(self):
        # Compute Layers dependency graph, as both adjacency and incidence list.
        layerdeps_adj = dict()
        layerdeps_inc = dict()

        for node in self.nodes:
            for layer in node.layers:
                if layer not in layerdeps_adj:
                    layerdeps_adj[layer] = set()

            for upper in node.registrations:
                for lower in node.registrations[upper]:
                    if upper not in layerdeps_inc:
                        layerdeps_inc[upper] = set()
                    if lower not in layerdeps_inc:
                        layerdeps_inc[lower] = set()
                    if upper not in layerdeps_adj:
                        layerdeps_adj[upper] = set()
                    if lower not in layerdeps_adj:
                        layerdeps_adj[lower] = set()
                    layerdeps_inc[upper].add(lower)
                    layerdeps_adj[lower].add(upper)

        # Kahn's algorithm below only needs per-node count of
        # incident edges, so we compute these counts from the
        # incidence list and drop the latter.
        layerdeps_inc_cnt = dict()
        for layer in layerdeps_inc:
            layerdeps_inc_cnt[layer] = len(layerdeps_inc[layer])
        del layerdeps_inc

        # Init layerdeps_inc_cnt for those Layers that do not
        # act as lower IPCPs nor upper IPCPs for registration
        # operations
        for node in self.nodes:
            for layer in node.layers:
                if layer not in layerdeps_inc_cnt:
                    layerdeps_inc_cnt[layer] = 0

        # Run Kahn's algorithm to compute topological
        # ordering on the Layers graph.
        frontier = set()
        self.layer_ordering = []
        for layer in layerdeps_inc_cnt:
            if layerdeps_inc_cnt[layer] == 0:
                frontier.add(layer)

        while len(frontier):
            cur = frontier.pop()
            self.layer_ordering.append(cur)
            for nxt in layerdeps_adj[cur]:
                layerdeps_inc_cnt[nxt] -= 1
                if layerdeps_inc_cnt[nxt] == 0:
                    frontier.add(nxt)
            layerdeps_adj[cur] = set()

        circular_set = [layer for layer in layerdeps_inc_cnt
                        if layerdeps_inc_cnt[layer] != 0]
        if len(circular_set):
            raise Exception("Fatal error: The specified Layers topology"
                            "has one or more"
                            "circular dependencies, involving the following"
                            " Layers: %s" % circular_set)

        logger.debug("Layer topological ordering: %s", self.layer_ordering)

    # Compute all the enrollments, to be called after compute_layer_ordering()
    def compute_enrollments(self):
        layer_graphs = dict()
        self.enrollments = []
        self.flows = []

        for layer in self.layer_ordering:
            neighsets = dict()
            layer_graphs[layer] = dict()
            first = None

            # For each N-1-Layer supporting this Layer, compute the set of nodes
            # that share such N-1-Layer. This set will be called the 'neighset' of
            # the N-1-Layer for the current Layer.

            for node in self.nodes:
                if layer in node.registrations:
                    layer_graphs[layer][node] = []  # init for later use
                    if first is None:  # pick any node for later use
                        first = node
                    for lower_layer in node.registrations[layer]:
                        if lower_layer not in neighsets:
                            neighsets[lower_layer] = []
                        neighsets[lower_layer].append(node)

            # Build the graph, represented as adjacency list
            for lower_layer in neighsets:
                # Each neighset corresponds to a complete (sub)graph.
                for node1 in neighsets[lower_layer]:
                    for node2 in neighsets[lower_layer]:
                        if node1 != node2:
                            layer_graphs[layer][node1].append((node2, lower_layer))

            self.enrollments.append([])
            self.flows.append([])

            if first is None:
                # This is a shim Layer, nothing to do
                continue

            er = []
            for node in layer_graphs[layer]:
                for edge in layer_graphs[layer][node]:
                    er.append("%s --[%s]--> %s" % (node.name,
                                                   edge[1].name,
                                                   edge[0].name))
            logger.debug("Layer graph for %s: %s", layer, ', '.join(er))

            # To generate the list of minimal enrollments
            # and minimal flows, we simulate it, using breadth-first
            # traversal from a graph-central root (shallow tree).
            root = _enrollment_root(layer_graphs[layer])

            enrolled = {root}
            frontier = {root}
            edges_covered = set()
            while len(frontier):
                cur = frontier.pop()
                for edge in layer_graphs[layer][cur]:
                    if edge[0] not in enrolled:
                        enrolled.add(edge[0])
                        enrollee = edge[0].get_ipcp_by_layer(layer)
                        assert(enrollee is not None)
                        enroller = cur.get_ipcp_by_layer(layer)
                        assert(enroller is not None)
                        edges_covered.add((enrollee, enroller))
                        self.enrollments[-1].append({'layer': layer,
                                                     'enrollee': enrollee,
                                                     'enroller': enroller,
                                                     'lower_layer': edge[1]})
                        self.flows[-1].append({'src': enrollee,
                                               'dst': enroller})
                        frontier.add(edge[0])
            if len(layer.members) != len(enrolled):
                raise Exception("Disconnected Layer found: %s" % (layer,))

            # In case of a full mesh enrollment or dt flows
            for cur in layer_graphs[layer]:
                for edge in layer_graphs[layer][cur]:
                    if cur.name < edge[0].name:
                        enrollee = cur.get_ipcp_by_layer(layer)
                        assert(enrollee is not None)
                        enroller = edge[0].get_ipcp_by_layer(layer)
                        assert(enroller is not None)
                        if ((enrollee, enroller) not in edges_covered and
                            (enroller, enrollee) not in edges_covered):
                            if self.enrollment_strategy == 'full-mesh':
                                self.enrollments[-1].append({'layer': layer,
                                                             'enrollee': enrollee,
                                                             'enroller': enroller,
                                                             'lower_layer': edge[1]})
                            if self.flows_strategy == 'full-mesh':
                                self.flows[-1].append({'src': enrollee,
                                                          'dst': enroller})
                            edges_covered.add((enrollee, enroller))

            if not (self.flows_strategy == 'minimal'
                    or self.flows_strategy == 'full-mesh') \
                    or not (self.enrollment_strategy == 'full-mesh'
                            or self.enrollment_strategy == 'minimal'):
                # This is a bug
                assert False

        log_string = "Enrollments:\n"
        for el in self.enrollments:
            for e in el:
                log_string += ("    [%s] %s --> %s through N-1-Layer %s\n"
                               % (e['layer'],
                                  e['enrollee'].name,
                                  e['enroller'].name,
                                  e['lower_layer']))
        logger.debug(log_string)

        log_string = "Flows:\n"
        for el in self.flows:
            for e in el:
                log_string += ("    %s --> %s \n"
                               % (e['src'].name,
                                  e['dst'].name))
        logger.debug(log_string)

    def compute_ipcps(self):
        # For each node, compute the required IPCP instances, and associated
        # registrations
        for node in self.nodes:
            node.ipcps = []
            # We want also the node.ipcps list to be generated in
            # topological ordering
            for layer in self.layer_ordering:
                if layer not in node.layers:
                    continue

                # Create an instance of the required IPCP class
                ipcp = layer.get_ipcp_class()(
                    name='%s.%s' % (layer.name, node.name),
                    node=node, layer=layer)

                if layer in node.registrations:
                    for lower in node.registrations[layer]:
                        ipcp.registrations.append(lower)

                node.ipcps.append(ipcp)
                layer.ipcps.append(ipcp)

    def compute_bootstrappers(self):
        for node in self.nodes:
            for ipcp in node.ipcps:
                ipcp.bootstrapper = True
                for el in self.enrollments:
                    for e in el:
                        if e['layer'] != ipcp.layer:
                            # Skip this Layer
                            break
                        if e['enrollee'] == ipcp:
                            ipcp.bootstrapper = False
                            # Exit the loops
                            break
                    if not ipcp.bootstrapper:
                        break

    def dump_ssh_info(self):
        f = open(os.path.join(tmp_dir, 'ssh_info'), 'w')

        for node in self.nodes:
            f.write("%s;%s;%s;%s;%s\n" % (node.name,
                                          self.testbed.username,
                                          node.ssh_config.hostname,
                                          node.ssh_config.port,
                                          node.ssh_config.proxy_server))
        f.close()

    # Examine the nodes and Layers, compute the registration and enrollment
    # order, the list of IPCPs to create, registrations, ...
    def generate(self):
        start = time.time()
        self.compute_layer_ordering()
        self.compute_ipcps()
        self.compute_enrollments()
        self.compute_bootstrappers()
        for node in self.nodes:
            logger.info("IPCPs for node %s: %s", node.name, node.ipcps)
        end = time.time()
        logger.info("Layer ordering computation took %.2f seconds", end - start)

    def install_prototype(self):
        """
        Installs the prototype on the nodes.
        """
        start = time.time()
        self._install_prototype()
        end = time.time()
        logger.info("Install took %.2f seconds", end - start)

    def set_startup_command(self, command):
        for node in self.nodes:
            node.startup_command = command

    def bootstrap_prototype(self):
        """
        Bootstraps the prototype on the nodes.
        """
        start = time.time()
        self._bootstrap_prototype()
        end = time.time()
        logger.info("Bootstrap took %.2f seconds", end - start)

    def update_prototype(self, branch=None):
        """
        Updates from branch and reinstalls and rebootstraps the prototype
        """
        start = time.time()
        self._terminate_prototype(force=True)
        self._update_prototype(branch)
        self._bootstrap_prototype()
        end = time.time()
        logger.info("Prototype update took %.2f seconds", end - start)

    def start_metrics_exporter(self, nodes=None, interval=100):
        """
        Start the metrics exporter for the prototype on nodes exporting data every interval ms
        """
        start = time.time()
        self._start_metrics_exporter(nodes, interval)
        end = time.time()
        logger.info("Starting metrics exporters took %.2f seconds", end - start)

    def stop_metrics_exporter(self, nodes=None):
        """
        Stop the metrics exporters for the prototype on nodes exporting data every interval ms
        """
        start = time.time()
        self._stop_metrics_exporter(nodes)
        end = time.time()
        logger.info("Stopping metrics exporters took %.2f seconds", end - start)

    @abc.abstractmethod
    def destroy_layer(self, layer):
        raise Exception('destroy_layer() method not implemented')

    @abc.abstractmethod
    def _install_prototype(self):
        raise Exception('install_prototype() method not implemented')

    @abc.abstractmethod
    def _bootstrap_prototype(self):
        raise Exception('bootstrap_prototype() method not implemented')

    @abc.abstractmethod
    def _update_prototype(self, branch):
        raise Exception('update_prototype() method not implemented')

    @abc.abstractmethod
    def prototype_name(self):
        raise Exception('prototype_name() method not implemented')

    @abc.abstractmethod
    def _terminate_prototype(self, force):
        raise Exception('terminate_prototype() method not implemented')

    @abc.abstractmethod
    def _start_metrics_exporter(self, nodes, interval):
        raise Exception('start_metrics_exporter() method not implemented')

    @abc.abstractmethod
    def _stop_metrics_exporter(self, nodes):
        raise Exception('start_metrics_exporter() method not implemented')

    @abc.abstractmethod
    def _set_phy_link_rate_to(self, src, dst, megabit):
        raise Exception('set_phy_k')

    def set_phy_link_rate_to(self, src, dst, megabit):
        """
        Sets the physical link rate in the experiment
        """
        self._set_phy_link_rate_to(src, dst, megabit)

    def set_phy_link_rate_between(self, src, dst, megabit):
        """
        Sets the physical link rate in the experiment
        """
        self._set_phy_link_rate_to(src, dst, megabit)
        self._set_phy_link_rate_to(dst, src, megabit)

    def swap_in(self):
        """
        Swap the experiment in on the testbed.
        """
        start = time.time()
        self.testbed.swap_in(self)
        self.dump_ssh_info()
        end = time.time()
        logger.info("Swap-in took %.2f seconds", end - start)

    def swap_out(self):
        """
        Swap the experiment out of the testbed.
        """
        start = time.time()
        for node in self.nodes:
            if node.ssh_config.client is not None:
                node.ssh_config.client.close()
            if node.ssh_config.proxy_client is not None:
                node.ssh_config.proxy_client.close()
        # Undo the testbed (testbed-specific)
        self.testbed.swap_out(self)
        end = time.time()
        logger.info("Swap-out took %.2f seconds", end - start)

    def terminate_prototype(self, force=False):
        """
        Terminate the prototype in the experiment.
        """
        self._terminate_prototype(force=force)

    def reboot_nodes(self):
        """
        Reboot all nodes in the experiment.
        """
        for node in self.nodes:
            node.execute_command('reboot', as_root=True)

    @abc.abstractmethod
    def export_layer_bandwidth(self, filename, layer):
        raise Exception('Export Layer bandwidth method not implemented')

    def export_connectivity_graph(self, filename):
        """
        Generate a PDF of the physical connectivity graph.

        :param filename: The output PDF filename.
        """
        try:
            import pydot

            colors = ['red', 'green', 'blue', 'orange', 'yellow']
            fcolors = ['black', 'black', 'white', 'black', 'black']

            gvizg = pydot.Dot(graph_type='graph')

            i = 0
            try:
                for node in self.nodes:
                    g_node = pydot.Node(node.name,
                                        label=node.name,
                                        style="filled",
                                        fillcolor=colors[i],
                                        fontcolor=fcolors[i])

                    gvizg.add_node(g_node)
                    i += 1
                    if i == len(colors):
                        i = 0
            except Exception as e:
                logger.error('Failed to create pydot Node: ' + str(e))
                return

            try:
                for layer in self.layer_ordering:
                    if layer.is_eth:
                        edge = pydot.Edge(layer.members[0].name,
                                          layer.members[1].name,
                                          label=layer.name,
                                          color='black')
                        gvizg.add_edge(edge)
            except Exception as e:
                logger.error('Failed to create pydot Edge: ' + str(e))
                return

            try:
                gvizg.write_pdf(filename)
                logger.info('Generated pdf of connectivity graph')
            except Exception as e:
                logger.error('Failed to write PDF: ' + str(e))

        except:
            logger.error('Warning: pydot module not installed, '
                         'cannot produce graph images')

    def export_layer_graph(self, filename, layer):
        """
        Generate a PDF of a Layer graph.

        :param filename: The output PDF filename.
        :param layer: The Layer to export.
        """
        try:
            import pydot

            colors = ['red', 'green', 'blue', 'orange', 'yellow']
            fcolors = ['black', 'black', 'white', 'black', 'black']

            gvizg = pydot.Dot(graph_type='digraph')

            i = 0
            try:
                for node in layer.members:
                    g_node = pydot.Node(node.name,
                                        label=node.name,
                                        style="filled",
                                        fillcolor=colors[i],
                                        fontcolor=fcolors[i])

                    gvizg.add_node(g_node)
                    i += 1
                    if i == len(colors):
                        i = 0
            except Exception as e:
                logger.error('Failed to create pydot Node: ' + str(e))
                return

            try:
                for enroll, dt in zip(self.enrollments,
                                      self.flows):
                    for e in enroll:
                        if e['layer'] is not layer:
                            continue

                        edge = pydot.Edge(e['enrollee'].node.name,
                                          e['enroller'].node.name,
                                          color='black')
                        gvizg.add_edge(edge)

                    for e in dt:
                        if e['src'].layer is not layer:
                            continue

                        edge = pydot.Edge(e['src'].node.name,
                                          e['dst'].node.name,
                                          color='red')
                        gvizg.add_edge(edge)
            except Exception as e:
                logger.error('Failed to create pydot Edge: ' + str(e))
                return

            try:
                gvizg.write_pdf(filename)
                logger.info('Generated PDF of Layer graph')
            except Exception as e:
                logger.error('Failed to write PDF: ' + str(e))

        except:
            logger.error('Warning: pydot module not installed, '
                         'cannot produce Layer graph images')


class Executor:
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def execute_command(self, node, command, as_root=False, time_out=3,
                        check=True):
        """Execute command on a node.

        :param node: The node to execute on.
        :param command: The command string.
        :param as_root: Execute as root.
        :param time_out: Timeout in seconds.
        :param check: If True, raise CommandError on non-zero exit.
        :returns: A CommandResult.
        """
        return

    def execute_commands(self, node, commands, as_root=False, time_out=3):
        for command in commands:
            self.execute_command(node, command, as_root, time_out)

    @abc.abstractmethod
    def copy_file(self, node, path, destination):
        return

    def copy_files(self, node, paths, destination):
        for path in paths:
            self.copy_file(node, path, destination)

    @abc.abstractmethod
    def fetch_file(self, node, path, destination, sudo=False):
        return

    def fetch_files(self, node, paths, destination, sudo=False):
        for path in paths:
            self.fetch_file(node, path, destination, sudo)
