#!/usr/bin/env python

"""
    Requirements:
        pydot
"""

import pvl.args
import pvl.hosts
from pvl.invoke import merge

import collections
import logging; log = logging.getLogger('pvl.hosts-graph')
import optparse

COLOR_VLANS = {
    1:      'grey',         # pvl-lan
    2:      'blue',         # pvl-lan2
    3:      'red',          # pvl-san
    4:      'green',        # pvl-veturi
    7:      'orange',       # pvl-ranssi
    8:      'yellow',       # pvl-mgmt
    10:     'brown',        # pvl-public
    100:    'navyblue',     # pvl-test
    103:    'red4',         # pvl-test-san
    104:    'red2',         # pvl-ganeti
    192:    'purple',       # paivola-services
    255:    'magenta',      # pvl-sonera
}

class ParseError (Exception) :
    def __init__ (self, file, line, msg) :
        self.file = file
        self.line = line
        self.msg = msg
  
    def __str__ (self) :
        return "{self.file}:{self.line}: {self.msg}".format(self=self)

def _parse_snmp_part (part) :
    if part.isdigit() :
        return int(part)
    else :
        return part

def _parse_snmp_attr (line) :
    for part in line.split() :
        yield _parse_snmp_part(part)

def _parse_snmp_value (line) :
    if '\t' in line :
        key, value = line.split('\t', 1)

        return { _parse_snmp_part(key): _parse_snmp_part(value) }

    else :
        return set((_parse_snmp_part(line), ))
    
def _load_snmp_data (options, file) :
    """
        Load a data dict generated by pvl.hosts-snmp from a file.

        Yields (host, attr, value)
    """

    host = None
    attr = None
    value = None
    
    for idx, line in enumerate(file, 1) :
        indent = 0

        while line.startswith('\t') :
            indent += 1
            line = line[1:]

        line = line.lstrip('\t').rstrip('\n')

        if indent == 0 :
            host = line
            attr = None
            value = None

        elif indent == 1 :
            attr = tuple(_parse_snmp_attr(line))
            value = None

            yield host, attr, None

        elif indent == 2 :
            if not attr :
                raise ParseError(file, line, "[%s] %s: value outside of attr" % (host, attr))
            
            value = _parse_snmp_value(line)

            yield host, attr, value

def load_snmp_data (options, file, hosts) :
    """
        Load snmp data as dict, from given file path, or stdin.
    """

    if file :
        file = open(file)
    else :
        file = sys.stdin

    root = { }

    for host_name, attr, value in _load_snmp_data(options, file) :
        host = hosts.get(host_name)
        
        if value :
            log.debug("[%s] %s: %s", host, ' '.join(str(a) for a in attr), value)
        else :
            log.debug("[%s] %s", host, ' '.join(str(a) for a in attr),)

        item = root.setdefault(host, { })
        
        for a in attr[:-1] :
            item = item.setdefault(a, {})

        a = attr[-1]
        
        if value is None :
            pass

        elif isinstance(value, set) :
            item.setdefault(a, set()).update(value)

        elif isinstance(value, dict) :
            item.setdefault(a, dict()).update(value)

        else :
            item[a] = value
            
    return root

def host_vlans (host, host_vlans) :
    """
        {vlan: { tagged/untagged: [port] } } -> (port, (untag, [tag])).
    """

    ports = set()
    vlans_untagged = { }
    vlans_tagged = collections.defaultdict(set)

    for vlan, vlan_attrs in host_vlans.iteritems() :
        for port in vlan_attrs.get('tagged', ()) :
            ports.add(port)
            vlans_tagged[port].add(vlan)
        
        for port in vlan_attrs.get('untagged', ()) :
            ports.add(port)
            vlans_untagged[port] = vlan
    
    for port in ports :
        untag = vlans_untagged.get(port)
        tagged = vlans_tagged.get(port, ())

        log.debug("%s: %s: untag=%s tag=%s", host, port, untag, tagged)
        
        yield port, (untag, tagged)

def build_graph (snmp, hosts) :
    """
        Combine given snmp data and { host: Host } into
            { node: label }
            { (remote, remote_port, local_port, local): (local_untag, tagged, remote_untag) }
    """

    nodes = { } # host: label
    links = { }

    hosts_by_lldp = { }
    
    # first scan: lldp hosts
    for host, host_attrs in snmp.iteritems() :
        lldp = host_attrs.get('lldp')

        if lldp :
            lldp_local = lldp['local']
            local = lldp_local['chassis']
            
            nodes[host] = host.location
            hosts_by_lldp[local] = host
    
    # second scan: lldp remotes
    for host, host_attrs in snmp.iteritems() :
        lldp = host_attrs.get('lldp')

        if not lldp :
            continue

        local = lldp['local']['chassis']
        local_node = host

        if 'vlan' in host_attrs :
            vlans = dict(host_vlans(host, host_attrs['vlan']))
        else :
            vlans = None
        
        for port, port_attrs in lldp.get('port', { }).iteritems() :
            local_port = port_attrs['local']['port']

            for remote, remote_attrs in port_attrs['remote'].iteritems() :
                if remote in hosts_by_lldp :
                    remote_node = hosts_by_lldp[remote]
                else :
                    remote_node = remote

                    # non-snmp lldp host
                    nodes[remote_node] = remote_attrs['sys_name']

                remote_port = remote_attrs['port']
                
                # local vlans
                if vlans :
                    port_vlans = vlans.get(port)
                else :
                    port_vlans = None

                if port_vlans :
                    local_untag, local_tagged = port_vlans
                
                # bidirectional mappings
                forward = (local_node, local_port, remote_port, remote_node)
                reverse = (remote_node, remote_port, local_port, local_node)

                if reverse not in links :
                    links[forward] = (local_untag, local_tagged, None)
                else :
                    remote_untag, remote_tagged, _ = links[reverse]

                    # merge
                    if remote_untag != local_untag :
                        log.warning("%s:%s untag %s <=> %s untag %s:%s",
                                host, local_port, local_untag,
                                remote_untag, remote_node, remote_port
                        )

                    if remote_tagged != local_tagged :
                        log.warning("%s:%s tagged %s <-> %s tagged %s:%s",
                                host, local_port, ':'.join(str(x) for x in sorted(local_tagged)),
                                ':'.join(str(x) for x in sorted(remote_tagged)), remote_node, remote_port
                        )

                    links[reverse] = (remote_untag, remote_tagged, local_untag)

    return nodes, links

class GraphVlans (object) :
    """
        Maintain vlan -> dot style/color mappings
    """

    SERIES = 'paired12'
    NONE = 'black'

    def __init__ (self, vlans=None) :
        if vlans :
            self.vlans = dict(vlans)
        else :
            self.vlans = { }
    
    def color (self, vlan) :
        if vlan in self.vlans :
            return self.vlans[vlan]
        
        # alloc
        color = '/{series}/{index}'.format(series=self.SERIES, index=len(self.vlans) + 1)

        self.vlans[vlan] = color

        return color

def dot_quote (value) :
    """
        Quote a dot value.
    """

    return '"{value}"'.format(value=value)

def dot (*line, **attrs) :
    """
        Build dot-syntax:
            *line {
                *line [**attrs];
            }
    """

    if line and attrs :
        return ''.join(('\t', ' '.join(str(x) for x in line), ' [',
            ', '.join('{name}="{value}"'.format(name=name, value=value) for name, value in attrs.iteritems()),
        ']', ';'))
    elif line :
        return ' '.join(line) + ' {'
    else :
        return '}'

def build_dot (options, nodes, links, type='digraph', vlans=None) :
    """
        Construct a dot description of the given node/links graph.
    """

    if vlans is True :
        vlans = { }

    yield dot(type, 'verkko')

    # defaults
    yield dot('graph',
            # XXX: breaks multi-edges?
            #splines     = 'true',

            sep             = '+25,25',
            overlap         = 'scalexy',

            # only applies to loops
            nodesep     = 0.5,
    )
    yield dot('edge',
        labeldistance   = 3.0,
        penwidth        = 2.0,
    )
    yield dot('node',
        fontsize        = 18,
    )
    
    # nodes
    for node, node_label in nodes.iteritems() :
        yield dot(dot_quote(node), label=node_label)

    # links
    for (local, local_port, remote_port, remote), (local_untag, tagged, remote_untag) in links.iteritems() :
        if vlans :
            head_color = vlans.color(local_untag) if local_untag else None
            tail_color = vlans.color(remote_untag) if remote_untag else None
            line_colors = [vlans.color(tag) for tag in sorted(tagged)]
        else :
            head_color = GraphVlans.NONE if local_untag else None
            tail_color = GraphVlans.NONE if remote_untag else None
            line_colors = []

        if head_color and tail_color :
            dir = 'both'
            colors = [head_color, tail_color] + line_colors
        elif head_color :
            dir = 'forward'
            colors = [head_color] + line_colors
        elif tail_color :
            dir = 'back'
            colors = [vlans.NONE, tail_color] + line_colors
        else :
            dir = 'none'
            colors = line_colors

        yield dot(dot_quote(local), '->', dot_quote(remote),
            taillabel   = local_port,
            headlabel   = remote_port,
            dir         = dir,

            fillcolor   = 'black',
            color       = ':'.join(colors),
        )

    yield dot()

def apply_dot (options, file, dot) :
    """
        Output dot file for given graphbits
    """

    for line in dot :
        file.write(line + '\n')

def main (argv) :
    """
        Graph network
    """

    parser = optparse.OptionParser(main.__doc__)
    parser.add_option_group(pvl.args.parser(parser))
    parser.add_option_group(pvl.hosts.optparser(parser))

    parser.add_option('--snmp-data', metavar='FILE', default=None,
            help="Load snmp data from FILE")

    parser.add_option('--graph-dot', metavar='FILE',
            help="Output .dot graph data to file")

    parser.add_option('--graph-vlans', action='store_true', dest='graph_vlans', 
            help="Graph all VLANs")

    parser.add_option('--no-vlans', action='store_false', dest='graph_vlans',
            help="Do not color VLANs")

    # input
    options, args = parser.parse_args(argv[1:])
    pvl.args.apply(options)
    
    # load hosts for correlation
    hosts = dict((str(host), host) for host in pvl.hosts.apply(options, args))

    # load raw snmp data
    snmp = load_snmp_data(options, options.snmp_data, hosts)

    # process data into graph
    nodes, links = build_graph(snmp, hosts)
    
    # process graph into dot
    if options.graph_vlans is False :
        graph_vlans = None
    else :
        graph_vlans = GraphVlans()

    if options.graph_dot :
        # process to dot
        dot = build_dot(options, nodes, links, vlans=graph_vlans)
        
        # write out
        apply_dot(options, open(options.graph_dot, 'w'), dot)

    return 0

if __name__ == '__main__':
    pvl.args.main(main)