#!/usr/bin/env python3
import datetime
from itertools import groupby
from operator import itemgetter
from pprint import pformat
from argparse import ArgumentParser, RawTextHelpFormatter
from execo import logger, Process, format_date
from execo.log import style
from execo_g5k import get_g5k_sites, get_oar_job_nodes, get_oargrid_job_nodes, \
    get_cluster_site, get_host_site, get_host_cluster, get_g5k_clusters, \
    get_g5k_hosts, get_host_shortname
from execo_g5k.topology import g5k_graph, treemap
from xml.dom import minidom
from xml.etree.ElementTree import Element, SubElement, tostring, dump
from decimal import Decimal
from networkx.algorithms.shortest_paths.generic import shortest_path
from execo_g5k.api_utils import get_cluster_network_equipments,\
    get_site_clusters, get_cluster_hosts, group_hosts


# Initializing topo5k
suffix = '.grid5000.fr'
sites = []
clusters = []
hosts = []
prog = 'topo5k'
version = '0.1'
description = 'This tool provides an access to the network topology of ' + \
    'the Grid\'5000 platform. It has two output format: \n' + \
    style.host('treemap') + ' = generate a map with networkx\n' + \
    style.host('simgrid') + ' = produce a SimGrid platform file (v4)'
epilog = ''


def main():
    """ """
    args = set_options()
    # Checking that mode given is valid
    if args.mode not in ['treemap', 'simgrid']:
        logger.error('Invalid mode selected: %s, \ntype %s for help',
                     style.emph(args.mode), style.emph(prog + ' -h'))
        exit()
    set_log_level(args)
    elements = define_elements(args)
    outfile = set_outfile(args, elements)

    gr = g5k_graph(elements)
    logger.debug('Nodes \n%s', pformat(sorted(gr.nodes())))
    logger.debug('Edges \n%s', pformat(sorted(gr.edges())))

    if args.mode == 'treemap':
        logger.info('Saving image file to %s', style.emph(outfile))
        plot = treemap(gr, compact=args.compact)
        plot.savefig(outfile, dpi=args.dpi)
        plot.clf()
    elif args.mode == 'simgrid':
        simgrid(gr, outfile=outfile)
        if args.check_simgrid:
            check = check_simgrid(outfile)
            if check:
                logger.info('Simgrid generated file is working')
            else:
                logger.info('Error in generated file')


def set_options():
    """ """
    parser = ArgumentParser(prog=prog,
                            description=description,
                            epilog=epilog,
                            formatter_class=RawTextHelpFormatter,
                            add_help=False)
    optinout = parser.add_argument_group(style.host("General options"),
                                         "Define mode and controls I/O.")
    optinout.add_argument("-h", "--help",
                          action="help",
                          help="show this help message and exit")
    optinout.add_argument('--version',
                          action='version',
                          version='%(prog)s ' + version)
    optinout.add_argument("-m", "--mode",
                          dest="mode",
                          default='treemap',
                          help="Setup the mode: treemap or simgrid " +
                          "\ndefault = %(default)s")
    optinout.add_argument("-o", "--outfile",
                          dest="outfile",
                          help="Define the outfile name")
    optinout.add_argument('--compact',
                          dest='compact',
                          action="store_true",
                          help="Use cluster description instead of hosts")

    optio = optinout.add_mutually_exclusive_group()
    optio.add_argument("-q", "--quiet",
                       dest="quiet",
                       action="store_true",
                       default=False,
                       help="Run without printing anything")
    optio.add_argument("-v", "--verbose",
                       dest="verbose",
                       action="store_true",
                       default=False,
                       help="Run in verbose mode")

    resources = parser.add_mutually_exclusive_group()
    resources.add_argument("-r", "--resources",
                           dest="resources",
                           default="grid5000",
                           help="Comma separated list of Grid'5000 " +
                           "elements (grid5000 or sites )")
    resources.add_argument("-j", "--job-id",
                           dest="job_id",
                           default=None,
                           help="oar_job_id (site:XXXXX) or oargrid_job_id")

    optmap = parser.add_argument_group(style.host("TREEMAP options"))
    optmap.add_argument('--graphviz-program',
                        dest='graphviz',
                        default='neato',
                        help='Graphviz program used for position')
    optmap.add_argument('--dpi',
                        type=int,
                        dest="dpi",
                        default=300,
                        help="DPI of the image")
    optmap.add_argument('--all-nodes',
                        action="store_true",
                        dest="all_nodes",
                        help="Show all nodes labels")

    optsimgrid = parser.add_argument_group(style.host("SIMGRID options"))
    optsimgrid.add_argument('--check-simgrid',
                            dest='check_simgrid',
                            action="store_true",
                            help="Execute a masterslave example on the " +
                            "generated platform file (requires SimGrid)")

    return parser.parse_args()


def set_log_level(args):
    if args.verbose:
        logger.setLevel('DEBUG')
    elif args.quiet:
        logger.setLevel('WARN')
    else:
        logger.setLevel('INFO')
    logger.info('%s', style.log_header('-- Grid\'5000 network topology --'))


def define_elements(args):
    """ """
    elements = []
    if args.job_id:
        logger.info('Waiting for job start ...')
        if ':' in args.job_id:
            for job in args.job_id.split(','):
                site, job_id = job.split(':')
                elements += map(lambda h: get_host_shortname(h.address),
                                get_oar_job_nodes(int(job_id), site))
        else:
            elements = map(lambda h: get_host_shortname(h.address),
                           get_oargrid_job_nodes(int(args.job_id)))
    else:
        if 'grid5000' in args.resources.split(','):
            elements = get_g5k_sites()
        else:
            for element in args.resources.split(','):
                if element in get_g5k_sites() + get_g5k_clusters() + \
                        get_g5k_hosts():
                    elements.append(element)
                else:
                    logger.warning('%s is not a valid element', element)

    logger.info('topology of %s', ' '.join(map(lambda x: style.emph(x),
                                               elements)))
    return elements


def set_outfile(args, elements):
    """Define the name of the output file"""
    if args.outfile is None:
        outfile = ''
        if args.job_id:
            if ':' in args.job_id:
                for job in args.job_id.split(','):
                    site, job_id = job.split(':')
                    outfile += site + '_' + job_id + '_'
            else:
                outfile = args.job_id + '.png'
        else:
            hosts = map(get_host_shortname, filter(lambda x: x in get_g5k_hosts(), elements))
            clusters = list(set([ get_host_cluster[e] for e in elements if get_host_cluster(e) in get_g5k_clusters() ]))
            sites = list(set([ get_cluster_site(c) for c in clusters ] + [ e for e in elements if e in get_g5k_sites() ]))

            if len(sites) == len(get_g5k_sites()):
                outfile = 'grid5000-' + str(datetime.date.today())
            else:
                outfile = '_'.join(sorted(sites)) + '-' + str(datetime.date.today())

        if args.mode == 'treemap':
            outfile += '.png'
        elif args.mode == 'simgrid':
            outfile += '.xml'
    else:
        outfile = args.outfile
    return outfile

def _get_site_router(gr, site):
    for n in gr.nodes:
        if n.endswith(f'.{site}') and gr.nodes[n]['kind'] == 'router':
            return n
    return None

def simgrid(gr, compact=True, outfile=None,
            tool_signature='Generated using %s %s' % (prog, version)):
    """ """
    def add_AS(elem, AS_id, routing="Full"):
        """Add an AS Element to an element"""
        logger.debug('Adding %s to %s', AS_id, elem)
        return SubElement(elem, 'AS', attrib={'id': AS_id, 'routing': routing})

    def add_site(elem, site, gr):
        """ """
        AS_site = add_AS(elem, 'AS_' + site)
        clusters = get_site_clusters(site)
        n_cluster = len(clusters)
        if n_cluster == 1:
            cluster = clusters[0]
            if len(_get_cluster_lcs(cluster, gr)) == 1:
                single_cluster_single_equip(cluster, AS_site, gr)
            else:
                add_router(AS_site, _get_site_router(gr, site))
                single_cluster_multi_equip(cluster, AS_site, gr)
                add_site_AS_routes(AS_site)
        else:
            add_router(AS_site, _get_site_router(gr, site))
            indep_cl, dep_cl = _clusters_dependance(clusters, gr)
            for c in indep_cl:
                if len(_get_cluster_lcs(c, gr)) == 1:
                    single_cluster_single_equip(c, AS_site, gr)
                else:
                    single_cluster_multi_equip(c, AS_site, gr)
            if len(dep_cl):
                cabinets(dep_cl, AS_site, gr)
            add_site_AS_routes(AS_site)

    def single_cluster_single_equip(cluster, AS_site, gr):
        """ """
        # retrieving first host data
        h = get_cluster_hosts(cluster)[0]
        power, cores = _flops_compact(gr.nodes[h]['power']), str(gr.nodes[h]['cores'])
        lc = next(iter(gr.neighbors(h)))
        # retrieving equip data
        bb_bw = _bw_compact(gr.nodes[lc]['backplane'])
        neighbors = gr.neighbors(lc)
        hosts = [ h for h in neighbors if gr.nodes[h]['kind'] == 'node' ]
        radical = _get_radical_list(hosts)
        # retrieving link data
        bw = _bw_compact([e for e in gr.edges(h, data=True) if lc in e][0][2]['bandwidth'])

        SubElement(AS_site, 'cluster',
                   attrib={'id': cluster,
                           'radical': radical,
                           'prefix': cluster + '-',
                           'suffix': '.grid5000.fr',
                           'speed': power,
                           'core': cores,
                           'bw': bw,
                           'bb_bw': bb_bw,
                           'lat': '0',
                           'router_id': lc})

    def single_cluster_multi_equip(cluster, AS_site, gr):
        """ """
        # Retrieving host power and cores
        h = get_cluster_hosts(cluster)[0]
        power, cores = _flops_compact(gr.nodes[h]['power']), str(gr.nodes[h]['cores'])
        # Retrieving equipments info
        lcs = {lc: _bw_compact(gr.nodes[lc]['backplane'])
               for lc in _get_cluster_lcs(cluster, gr)}
        i = 0
        for lc, bb_bw in lcs.items():
            hosts = [ h for h in gr.neighbors(lc) if gr.nodes[h]['kind'] == 'node' ]
            radical = _get_radical_list(hosts)
            bw = _bw_compact([e for e in gr.edges(next(iter(hosts)), data=True) if lc in e][0][2]['bandwidth'])
            SubElement(AS_site, 'cluster',
                       attrib={'id': cluster + '_' + str(i),
                               'radical': radical,
                               'prefix': cluster + '-',
                               'suffix': '.grid5000.fr',
                               'speed': power,
                               'core': cores,
                               'bw': bw,
                               'bb_bw': bb_bw,
                               'lat': '0',
                               'router_id': lc})
            i += 1

    def cabinets(dep_cl, AS_site, gr):
        """ """
        lcs = {}
        for c in dep_cl:
            for h in get_cluster_hosts(c):
                lc = next(iter(gr.neighbors(h)))
                if lc not in lcs:
                    lcs[lc] = _bw_compact(gr.nodes[lc]['backplane'])
        for lc, bb_bw in lcs.items():
            cabinets_as = add_AS(AS_site, 'AS_' + lc, 'Cluster')
            SubElement(cabinets_as, 'router', attrib={'id': lc})
            SubElement(cabinets_as, 'backbone', attrib={'id': 'backbone_' + lc,
                                                        'bandwidth': bb_bw,
                                                        'latency': str(0)})
            hosts = [ h for h in gr.neighbors(lc) if gr.nodes[h]['kind'] == 'node' ]
            grouped_hosts = group_hosts(hosts)
            for c, cl_h in grouped_hosts[AS_site.get('id').split('_')[1]].items():
                power, cores = _flops_compact(gr.nodes[cl_h[0]]['power']), str(gr.nodes[cl_h[0]]['cores'])
                bw = _bw_compact([e for e in gr.edges(cl_h[0], data=True) if lc in e][0][2]['bandwidth'])
                radical = _get_radical_list(cl_h)
                SubElement(cabinets_as, 'cabinet', 
                           attrib={'id': c,
                                   'radical': radical,
                                   'prefix': c + '-',
                                   'suffix': '.grid5000.fr',
                                   'speed': power,
                                 # 'core': cores,
                                   'bw': bw,
                                   'lat': '0'})

    def add_router(elem, router_id):
        """ """
        router_as = add_AS(elem, 'AS_' + router_id)
        SubElement(router_as, 'router', attrib={'id': router_id})
        return router_as

    def add_backbone(elem, sites):
        """ """
        raise NotImplementedError

    def add_site_AS_routes(AS_site):
        # Adding the ASroutes
        ASroutes = []
        router_element = AS_site.find('.//router')
        parent_router = router_element.get('id')
        AS_elements = AS_site.findall('./cluster') + AS_site.findall("./AS[@routing='Cluster']") 
        for AS in AS_elements:
            as1_id = AS.get('id')
            if AS.tag == 'cluster':
                router1 = AS.get('router_id')
            else:
                router1 = AS.find('./router').get('id')
            ASroutes.append([{'src': 'AS_' + parent_router,
                              'dst': as1_id,
                              'gw_src': parent_router,
                              'gw_dst': router1}, []])
            path = shortest_path(gr, router1, parent_router)

            links = []
            link_bw = None
            for j in range(len(path[:-1])):
                link_bw = _bw_compact(next(iter(gr[path[j]][path[j + 1]].values()))['bandwidth'])
                links.append({'id': path[j+1] + '_limiter',
                              'bandwidth': link_bw,
                              'latency': '0'})
            ASroutes[-1][1] = links
            for AS2 in AS_elements:
                as2_id = AS2.get('id')
                if as2_id != as1_id:
                    if AS2.tag == 'cluster':
                        router2 = AS2.get('router_id')
                    else:
                        router2 = AS2.find('./router').get('id')
                    ASroutes.append([{'src': as1_id,
                                      'dst': as2_id,
                                      'gw_src': router1,
                                      'gw_dst': router2},
                                     [{'id': parent_router + '_limiter',
                                       'bandwidth': link_bw,
                                       'latency': '0'}]])
        for _, links in ASroutes:
            for l in links:
                if AS_site.find(".//link[@id='" + l['id'] + "']") is None:
                    SubElement(AS_site, 'link', attrib=l)
        for asroute, links in ASroutes:
            if AS_site.find(".//ASroute[@src='" + asroute['dst'] + "']"
                            + "[@dst='" + asroute['src'] + "']") is None and \
               AS_site.find(".//ASroute[@src='" + asroute['src'] + "']"
                            + "[@dst='" + asroute['dst'] + "']") is None:
                route = SubElement(AS_site, 'ASroute', attrib=asroute)
                for l in links:
                    SubElement(route, 'link_ctn', id=l['id'])

    def _get_cluster_lcs(cluster, gr):
        """ """
        lcs = []
        for h in get_cluster_hosts(cluster):
            if h in gr.nodes():
                lc = next(iter(gr.neighbors(h)))
                if lc not in lcs:
                    lcs.append(lc)

        return lcs

    def _clusters_dependance(clusters, gr):
        """Check if the site clusters share a linecard"""
        lcs = {}
        for c in clusters:
            lcs[c] = _get_cluster_lcs(c, gr)

        dep_clusters = []
        indep_clusters = []
        for c1, eq1 in lcs.items():
            is_indep = True
            for c2, eq2 in lcs.items():
                if c2 != c1:
                    if len(set(eq1).intersection(set(eq2))):
                        is_indep = False
            if is_indep:
                indep_clusters.append(c1)
            else:
                dep_clusters.append(c1)

        return indep_clusters, dep_clusters

    def _subs_E(string):
        """ """
        conv = {0: '', 3: 'k', 6: 'M', 9: 'G', 12: 'T'}
        for pat, sub in conv.items():
            real_pat = 'E+' + "%d" % pat
            if real_pat in string:
                return string.replace(real_pat, sub)

    def _flops_compact(flops):
        """ """
        sci_not = Decimal(flops).normalize().to_eng_string()
        return _subs_E(sci_not) + 'f'

    def _bw_compact(bw):
        """ """
        if not bw:
            sci_not = "1E+12"
            logger.warning('no bw information, using %s', sci_not)
        else:
            sci_not = Decimal(bw / 8).normalize().to_eng_string()
        return _subs_E(sci_not) + 'Bps'

    def _get_radical_list(hosts):
        """ """
        radical = ''
        radical_list = sorted(map(lambda x: int(x.split('-')[1]), hosts))
        for k, g in groupby(enumerate(radical_list),
                            lambda t: t[0] - t[1]):
            radical_range = list(map(itemgetter(1), g))
            if len(radical_range) > 1:
                radical += str(min(radical_range)) + '-' + \
                    str(max(radical_range))
            else:
                radical += str(radical_range[0])
            radical += ','
        return radical[:-1]

    platform = Element('platform', attrib={'version': '4'})

    sites = gr.get_sites()
    # print gr.nodes(), filter(lambda x: x[2]['active'], gr.edges(data=True))
    if len(sites) == 1:
        add_site(platform, sites[0], gr)
    else:
        main_as = add_AS(platform, 'AS_grid5000')
        for site in sites:
            add_site(main_as, site, gr)
        add_backbone(main_as, sites)

    logger.info('Saving file to %s', style.emph(outfile))
    f = open(outfile, 'w')
    f.write('<?xml version="1.0" encoding="UTF-8"?>\n' +
            '<!DOCTYPE platform SYSTEM ' +
            '"http://simgrid.gforge.inria.fr/simgrid/simgrid.dtd">\n' +
            '<!-- ' + tool_signature + '\n     ' +
            'API commit ' + gr.graph['api_commit'] +
            '\n     ' + format_date(gr.graph['date']) + ' -->\n' +
            prettify(platform))
    f.close()


def check_simgrid(fname):
    """ """
    cmd = './examples/masterslave/masterslave_arg ' + \
        fname + ' 10000 1000'
    logger.info('Executing %s with generated platform file',
                style.command(cmd))
    Process('cd examples && tar -xzf masterslave.tgz && cd masterslave &&'
            'cmake . && make masterslave_arg',
            shell=True).run()
    test = Process(cmd).run()
    Process('rm -rf examples/masterslave').run()

    return test.ok


def prettify(elem):
    """Return a pretty-printed XML string for the Element"""
    rough_string = tostring(elem, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    return reparsed.toprettyxml(indent="  ").replace('<?xml version="1.0" ?>\n',
                                                     '')

if __name__ == "__main__":
    main()
