#!/usr/bin/env python
import datetime
from pprint import pprint
from argparse import ArgumentParser, RawTextHelpFormatter
from execo import logger, Process
from execo.log import style
from execo_g5k import get_g5k_sites, get_oar_job_nodes, get_oargrid_job_nodes, \
    get_host_site
from execo_g5k.topology import g5k_graph, treemap, hadoop, simgrid

# Initializing topo5k 
suffix = '.grid5000.fr'
sites = []
clusters = []
nodes = []

# Options definition
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' + \
    style.host('hadoop') + ' = return the list hosts grouped by switch '
epilog = ''

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 positions computations ')
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)")

# Parsing options
args = parser.parse_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 --'))

# Checking that mode given is valid
if args.mode not in ['treemap', 'simgrid', 'hadoop']:
    logger.error('Invalid mode selected: %s, \ntype %s for help',
                 style.emph(args.mode), style.emph(prog + ' -h'))
    exit()

# Defining sites and hosts
sites = []
hosts = []
if args.job_id:
    if ':' in args.job_id:
        tmp_outfile = ''
        for job in args.job_id.split(','):
            site, job_id = job.split(':')
            sites.append(site)
            hosts += get_oar_job_nodes(int(job_id), site)
            tmp_outfile += site + '_' + job_id + '_'
        if args.outfile is None:
            args.outfile = tmp_outfile[:-1] + '.png'
    else:
        hosts = get_oargrid_job_nodes(int(args.job_id))
        for node in nodes:
            site = get_host_site(node)
            if site not in sites:
                sites.append(site)
        if args.outfile is None:
            args.outfile = args.job_id + '.png'
else:
    if 'grid5000' in args.resources.split(','):
        sites = get_g5k_sites()
    else:
        for element in args.resources.split(','):
            if element in get_g5k_sites():
                sites.append(element)
            else:
                logger.warning('%S is not a valid site', element)
    if not args.outfile:
        if 'grid5000' in args.resources.split(','):
            args.outfile = 'grid5000-' + str(datetime.date.today())
        else:
            args.outfile = '_'.join(sites) + '-' + str(datetime.date.today())
        if args.mode == 'treemap':
            args.outfile += '.png'
        elif args.mode == 'simgrid':
            args.outfile += '.xml'


# Create the networkx graph object
if len(sites) == 0:
    logger.error('Wrong resources option given %s', args.resources)
    exit()
else:
    logger.info('topology of %s',
                ' '.join(map(lambda x: style.emph(x), sites)))
    if len(hosts) == 0:
        gr = g5k_graph(sites)
    else:
        gr = g5k_graph()
        for host in hosts:
            gr.add_host(host)


# Saving the graph object to image or XML
if args.mode == 'treemap':
    logger.info('Saving image file to %s', style.emph(args.outfile))
    plot = treemap(gr, compact=args.compact)
    plot.savefig(args.outfile, dpi=args.dpi)

elif args.mode == 'hadoop':
    hadoop(gr, hosts)

elif args.mode == 'simgrid':
    simgrid(gr, outfile=args.outfile, compact=args.compact, routing='Full')
    if args.check_simgrid:
        logger.info('Checking generated platform file')
        extract = Process('cd examples && tar -xzf masterslave.tgz')
        extract.shell = True
        extract.run()
        test = Process('./examples/masterslave/masterslave_arg ' +
                       args.outfile + ' 10 5').run()
        remove = Process('rm -rf examples/masterslave').run()
        if not test.ok:
            print test.stdout
            print test.stderr
        else:
            logger.info(style.command(test.cmd) + ' successfully executed.')

