#!/usr/bin/env python3
__license__ = """
Copyright 2009-2020 National Technology and Engineering Solutions of Sandia,
LLC (NTESS).  Under the terms of Contract DE-NA-0003525, the U.S.  Government 
retains certain rights in this software.

Sandia National Laboratories is a multimission laboratory managed and operated
by National Technology and Engineering Solutions of Sandia, LLC., a wholly 
owned subsidiary of Honeywell International, Inc., for the U.S. Department of 
Energy's National Nuclear Security Administration under contract DE-NA0003525.

Copyright (c) 2009-2020, NTESS

All rights reserved.

Redistribution and use in source and binary forms, with or without modification, 
are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.

    * Redistributions in binary form must reproduce the above
      copyright notice, this list of conditions and the following
      disclaimer in the documentation and/or other materials provided
      with the distribution.

    * Neither the name of the copyright holder nor the names of its
      contributors may be used to endorse or promote products derived
      from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

Questions? Contact <EMAIL>
"""

'''
 Loads a snapshot and parses parameters out of the fabric.
 Snapshots should be generated with the following
 ```bash
 opareport -o snapshot --route > snapshot-$(date -Iminutes).xml
 ```

Terminology
- LID (Local ID)
  - Represented as hex, but converted into integers in this script
  - Nodes will have one LID, and switches will have one per port
- GUID (Global Unique ID)
  - Represented as hex, but converted into integers in this script
- Name
  - A string identifier, similar to a hostname.
  - Switch names tend to be "{hostname}{GUID}"
  - Node names tend to be "{hostname} {Interface}" (note the space)
  - In every example I've seen, Names have a unique mapping to GUIDs
- Link
  - Snapshot enumerates connections between fabric nodes
  - Each consist of two (GUID, Port) pairs
'''


import argparse, sys, collections
import xml.etree.ElementTree as ET
from collections import defaultdict


def parse_args():
    parser = argparse.ArgumentParser(description='Parses reports generated by "opareport -o snapshot --route" into a JSON or YAML dictionary. Each enabled flag will populate out a key')
    parser.add_argument('-r', '--routing_table', action='store_true', help='Switch routing tables: GUID -> (lid -> port)')
    parser.add_argument('--named_routing_table', action='store_true', help='Switch routing tables with Names: Name -> (Name -> port)')
    parser.add_argument('-l', '--link_table', action='store_true', help='Fabric link table: GUID -> Port -> (GUID, Port)')
    parser.add_argument('-n', '--name_guid_map', action='store_true', help='Name to GUID map: Name -> GUID')
    parser.add_argument('-g', '--guid_name_map', action='store_true', help='GUID to Name map: GUID -> Name')
    parser.add_argument('-N', '--guid_lid_map', action='store_true', help='Node to LID map: GUID -> LID')
    parser.add_argument('-L', '--lid_guid_map', action='store_true', help='LID to Node map: LID -> GUID')
    parser.add_argument('-s', '--shorten_name', action='store_true', help='Shortens string names')
    parser.add_argument('--switch_links', metavar='SWITCH_NAME', help='Print switch links')
    parser.add_argument('--node_links', metavar='NODE_NAME', help='Print node links')
    parser.add_argument('-a', '--all', action='store_true', help='The kitchen sink!')
    parser.add_argument('-o', choices={'json', 'yaml'}, default='json', help='Output format')
    parser.add_argument('FILE', help='opareport xml file name.')

    return parser.parse_args()


def parse_xml_file(fname):
    return ET.parse(fname).getroot()


def parse_hex(lid):
    return int(lid, 0)


def get_guid(xml_element):
    return parse_hex(xml_element.attrib['id'])


def xml_is_switch(xml_element):
    return xml_element.find('./NodeType_Int').text == '2'


def get_nodes(xml_root):
    for n in xml_root.findall('./Nodes/Node'):
        if not xml_is_switch(n):
            yield n


def get_switches(xml_root):
    for sw in xml_root.findall('./Nodes/Node'):
        if xml_is_switch(sw):
            yield sw


def get_links(xml_root):
    for link in xml_root.findall('./Links/Link'):
        yield link


def gen_link_table(xml_root):
    '''guid -> port -> (guid, port)'''
    link_table_ = defaultdict(dict)
    for l in get_links(xml_root):
        from_guid = parse_hex(l.find('./From/NodeGUID').text)
        from_port = int(l.find('./From/PortNum').text)
        to_guid = parse_hex(l.find('./To/NodeGUID').text)
        to_port = int(l.find('./To/PortNum').text)

        # Make the table bidirectional
        link_table_[from_guid][from_port] = (to_guid, to_port)
        link_table_[to_guid][to_port] = (from_guid, from_port)
    return dict(link_table_) # convert to dict to play more nicely with the serializer


def get_link(guid, port):
    '''link table wrapper function for simplicity'''
    return link_table[guid][port] if key in link_table else False


def gen_routing_tables(xml_root):
    '''Switch GUID -> (lid -> port) Map'''
    switch_lid_port_map = {}
    for switch in get_switches(xml_root):
        guid = get_guid(switch)
        switch_lid_port_map[guid] = {parse_hex(entry.attrib['LID']): int(entry.text)
                                     for entry in switch.findall('./SwitchData/LinearFDB/')}
    return switch_lid_port_map


def iter_node_lids(xml_root):
    '''(node GUID, LID) generator'''
    for node in get_nodes(xml_root):
        guid = get_guid(node)
        lid = node.find('./PortInfo/LID').text
        yield (guid, parse_hex(lid))


def gen_node_lid_map(xml_root):
    '''node GUID -> LID'''
    return {n: l for n, l in iter_node_lids(xml_root)}


def gen_lid_node_map(xml_root):
    '''LID -> node GUID'''
    return {l: n for n, l in iter_node_lids(xml_root)}


def iter_node_name_guids(xml_root):
    '''(name, GUID) generator'''
    for n in xml_root.findall('./Nodes/Node'):
        name = n.find('./NodeDesc').text
        guid = get_guid(n)
        yield (name, guid)


def parse_if(true, name):
    '''Parses a node name when the first argument is true'''
    return name.split()[0] if true else name


def gen_guid_name_map(xml_root, parse_name=False):
    '''GUID -> name'''
    return {g: parse_if(parse_name, n) for n, g in iter_node_name_guids(xml_root)}


def gen_name_guid_map(xml_root, parse_name=False):
    '''name -> GUID'''
    return {parse_if(parse_name, n): g for n, g in iter_node_name_guids(xml_root)}


def switch_port_for_lid(switch_id, target_lid):
    '''For a given named switch and target LID, find the port to route to'''
    return routing_table[switch_id][target_lid]


def sst_map(name_prefix, guid_name_map, link_table):
    '''Generates a dictionary that can be parsed by SST'''
    name_map = {}
    unmapped = set()

    for guid, name in guid_name_map.items():
        # Determine whether we are parsing a node or switch
        if name_prefix not in name:
            unmapped.add(name)
            continue

        cur_map = name_map[name] = {'outports': {}}
        for (outport, (guid, inport)) in link_table[guid].items():
            cur_map['outports'][str(outport - 1)] = {
                'destination': guid_name_map[guid],
                'inport': inport - 1}
    return (name_map, unmapped)

def named_routing_tables(routing_table, guid_name_map, lid_guid_map):
    '''Generates a routing table with names instead of GUIDs'''
    new_rt = {}

    for (guid, lid_port_map) in routing_table.items():
        cur_id = new_rt[guid_name_map[guid]] = {'routes': {}}
        for (lid, port) in lid_port_map.items():
            if lid in lid_guid_map: # Switches don't have LIDs
                cur_id['routes'][guid_name_map[lid_guid_map[lid]]] = port - 1
    return new_rt

# https://gist.github.com/angstwad/bf22d1822c38a92ec0a9#gistcomment-1986197
def dict_merge(dct, merge_dct, add_keys=True):
    '''Merges two dictionaries into'''
    dct = dct.copy()
    if not add_keys:
        merge_dct = {
            k: merge_dct[k]
            for k in set(dct).intersection(set(merge_dct))
        }

    for k, v in merge_dct.items():
        if (k in dct and isinstance(dct[k], dict)
                and isinstance(merge_dct[k], collections.Mapping)):
            dct[k] = dict_merge(dct[k], merge_dct[k], add_keys=add_keys)
        else:
            dct[k] = merge_dct[k]

    return dct

if __name__ == '__main__':
    args = parse_args()
    xml_root = parse_xml_file(args.FILE)
    out_dict = {}

    if args.routing_table or args.all:
        out_dict['routing_table'] = gen_routing_tables(xml_root)
    if args.link_table or args.all:
        out_dict['link_table'] = gen_link_table(xml_root)
    if args.name_guid_map or args.all:
        out_dict['name_guid_map'] = gen_name_guid_map(xml_root, args.shorten_name)
    if args.guid_name_map or args.all:
        out_dict['guid_name_map'] = gen_guid_name_map(xml_root, args.shorten_name)
    if args.guid_lid_map or args.all:
        out_dict['guid_lid_map'] = gen_node_lid_map(xml_root)
    if args.lid_guid_map or args.all:
        out_dict['lid_guid_map'] = gen_lid_node_map(xml_root)
    if args.switch_links:
        (out_dict['switches'], unmapped_sw) = sst_map(args.switch_links,
                gen_guid_name_map(xml_root, args.shorten_name),
                gen_link_table(xml_root))
    if args.node_links:
        (out_dict['nodes'], unmapped_nd) = sst_map(args.node_links,
                gen_guid_name_map(xml_root, args.shorten_name),
                gen_link_table(xml_root))
    if args.named_routing_table:
        output = named_routing_tables(gen_routing_tables(xml_root),
                                     gen_guid_name_map(xml_root, args.shorten_name),
                                     gen_lid_node_map(xml_root))

        # If switch_links are used, we have to interleave the dicts b/c Jeremy says so
        out_dict['switches'] = dict_merge(out_dict['switches'], output) if args.switch_links else output

    # List fabric nodes that did not match either prefix
    if args.node_links and args.switch_links:
        unmapped = unmapped_sw.intersection(unmapped_nd)
        if len(unmapped) > 0 :
            print('Warning some fabric nodes unmapped:', file=sys.stderr)
            [print('  %s' % name, file=sys.stderr) for name in unmapped]

    if args.o == 'json':
        from json import dumps
        print(dumps(out_dict, sort_keys=True, indent=2))
    elif args.o == 'yaml':
        try:
            from yaml import safe_dump # safe_dump drops the annoying type tags
        except:
            print('ERROR: please install PyYAML for yaml output', file=sys.stderr)
            exit()
        print(safe_dump(out_dict, default_flow_style=False))
