#!/usr/bin/env python
#
#    vm5k: Automatic deployment of virtual machine on Grid'5000
#     Created by L. Pouilloux and M. Imbert (INRIA, 2013)
#
#    A great thanks to A. Lebre and J. Pastor for extensive testing.
#
#
import os
import sys
import argparse
from pprint import pformat
from time import strftime
from logging import INFO, DEBUG, WARN
from math import ceil
from xml.etree.ElementTree import fromstring, parse, dump
from execo import logger, Process, default_connection_params
from execo.log import style
from execo.time_utils import Timer, format_date, format_duration
from execo_g5k import oargridsub, oarsub
from execo_g5k.oar import get_oarsub_commandline
from execo_g5k.oargrid import get_oargridsub_commandline
from execo_g5k.planning import get_planning, compute_slots, show_resources, \
    find_free_slot, get_jobs_specs, distribute_hosts
# from execo_g5k.topology import g5k_graph, treemap
from execo_engine import copy_outputs
from vm5k import vm5k_deployment, define_vms, get_oar_job_vm5k_resources, \
    get_max_vms, get_oargrid_job_vm5k_resources, get_vms_slot, print_step

##############################################################################
# INITIALIZATION
##############################################################################
version = '0.6.1'
execution_time = {}
timer = Timer()


# Constants
deployment_tries = 1
blacklisted = ['sagittaire']
default_connection_params['user'] = 'root'

# Parsing options
prog = 'vm5k'
desc = 'A tool to deploy and configure nodes and virtual machines with ' + \
    style.log_header('Debian') + ' and ' + style.log_header('libvirt') + \
    '\non the ' + style.log_header('Grid5000') + ' platform. ' + \
    '.\nYou can use one of these options combinations:' + \
    '\n - ' + style.host('n_vm + oargrid_job_id') + \
    ' = use an existing reservation and specify number of VMs' + \
    '\n - ' + style.host('infile + oargrid_job_id') + \
    ' = use an existing reservation and specify vm placement XML file' + \
    '\n - ' + style.host('n_vm + walltime') + \
    ' = perform a reservation that has enough RAM' + \
    '\n - ' + style.host('infile + walltime') + \
    ' = perform a reservation according to the placement XML infile' + \
    '.\nBased on execo-2.2, ' + style.emph('http://execo.gforge.inria.fr/doc/')
epilog = style.host('Examples:') + '\nDeploy 100 VMs with the default ' + \
    'environnements for 3h ' + \
    style.command('\n %(prog)s -n 100 -w 3:00:00') + \
    '\nDocumentation can be found on ' + \
    style.emph('http://vm5k.readthedocs.org/en/latest/vm5k.html') + \
    '\nIssues/features requests can be reported to ' + \
    style.emph('https://github.com/lpouillo/vm5k')

parser = argparse.ArgumentParser(prog=prog, description=desc, epilog=epilog,
                                 formatter_class=argparse.RawTextHelpFormatter,
                                 add_help=False)
# Run options
run = parser.add_argument_group(style.host('Execution'),
                                "Manage how %(prog)s is executed")
run.add_argument("-h", "--help",
                 action="help",
                 help="show this help message and exit")
run.add_argument('--version',
                 action='version',
                 version='%(prog)s ' + version)

optio = run.add_mutually_exclusive_group()
optio.add_argument("-v", "--verbose",
                   action="store_true",
                   help='print debug messages')
optio.add_argument("-q", "--quiet",
                   action="store_true",
                   help='print only warning and error messages')
run.add_argument("-o", "--outdir",
                 dest="outdir",
                 default='vm5k_' + strftime("%Y%m%d_%H%M%S_%z"),
                 help='where to store the vm5k log files' +
                 "\ndefault=%(default)s")
run.add_argument("-p", "--program",
                 dest="program",
                 help='Launch a program at the end of the deployment')
run.add_argument("--plot",
                 dest='plot',
                 action="store_true",
                 help='draw a topological graph of the deployment')

# Reservation
mode = parser.add_argument_group(style.host("Mode"),
                                 "Define the mode of %(prog)s")
optnvm = mode.add_mutually_exclusive_group()
optnvm.add_argument('-n', '--n_vm',
                    dest='n_vm',
                    type=int,
                    help='number of virtual machines')
optnvm.add_argument('-i', '--infile',
                    dest="infile",
                    help='XML file describing the placement of VMs ' +
                    'on G5K sites and clusters')
optresa = mode.add_mutually_exclusive_group()
optresa.add_argument('-j', '--job_id',
                     dest='job_id',
                     help='use the hosts from a oargrid_job or a oar_job.')
optresa.add_argument('-w', '--walltime',
                     dest='walltime',
                     help='duration of your reservation')
mode.add_argument('-k', '--kavlan',
                  dest='kavlan',
                  action="store_true",
                  default=False,
                  help='Deploy the VMs in a KaVLAN')

# Hosts configuration
hosts = parser.add_argument_group(style.host('Physical hosts'),
                                  "Tune the physical hosts.")
hosts.add_argument('-r', '--resources',
                   dest='resources',
                   default='grid5000',
                   help='list of Grid\'5000 elements')
hosts.add_argument('-b', '--blacklisted',
                   dest='blacklisted',
                   help='list of Grid\'5000 elements to be blacklisted')
optenv = hosts.add_mutually_exclusive_group()
optenv.add_argument('-e', '--env_name',
                    dest='env_name',
                    help='Kadeploy environment name')
optenv.add_argument('-a', '--env_file',
                    dest='env_file',
                    help='path to the Kadeploy environment file')
optdeploy = hosts.add_mutually_exclusive_group()
optdeploy.add_argument('--forcedeploy',
                       action="store_true",
                       help='force the deployment of the hosts')
optdeploy.add_argument('--nodeploy',
                       action="store_true",
                       help='consider that hosts are already deployed and ' +
                       'configured')
hosts.add_argument('--no-packages-management',
                   dest='packages_management',
                   action='store_false',
                   help='disable package management')
hosts.add_argument('--other-packages',
                   dest='other_packages',
                   help='comma separated list of packages to be installed ' +
                   'on the hosts')


# VMs configuration
vms = parser.add_argument_group(style.host('Virtual machines'),
                                "Tune the virtual machines.")
vms.add_argument('-t', '--vm_template',
                 dest='vm_template',
                 help='XML string describing the virtual machine',
                 default='<vm mem="1024" hdd="10" n_cpu="1" cpuset="auto" />')
vms.add_argument('-f', '--vm_backing_file',
                 dest='vm_backing_file',
                 default='/grid5000/images/KVM/wheezy-x64-base.qcow2',
                 help='backing file for your virtual machines')
vms.add_argument('-l', '--vm_disk_location',
                 default='one',
                 dest='vm_disk_location',
                 help='Where to create the qcow2: one (default) or all)')
vms.add_argument('-d', '--vm_distribution',
                 dest='vm_distribution',
                 help='how to distribute the VMs round-robin (default) ' +
                 'n_by_hosts, random or concentrated')
vms.add_argument('--vm-clean-disks',
                 dest='vm_clean_disks',
                 action="store_true",
                 help='force to use a fresh copy of the vms backing_file')

# Services
service = parser.add_argument_group(style.host('Services'),
                                    "Deploy some services and hosts and vms")
service.add_argument('--aptcacher',
                     dest='aptcacher',
                     action="store_true",
                     help='configure aptcacher on hosts (servers) and ' +
                     'vms (clients)')

args = parser.parse_args()

# Create output directory and save execution log
try:
    os.mkdir(args.outdir)
except os.error:
    pass
copy_outputs(args.outdir + '/vm5k.log', args.outdir + '/vm5k.log')

# Set log level
out_deploy = False
if args.verbose:
    logger.setLevel(DEBUG)
    out_deploy = True
elif args.quiet:
    logger.setLevel(WARN)
else:
    logger.setLevel(INFO)

# Start message
print_step('VM5K: deployment of VMs on Grid\'5000 *')
logger.info('Version ' + version)

# Set demo values for n_vm and walltime
if args.n_vm is None and args.infile is None:
    logger.warning('No options: -n %s or -i %s, setting %s to %s',
                   style.emph('n_vm'), style.emph('infile'),
                   style.emph('n_vm'), style.emph(str(30)))
    args.n_vm = 30
if args.walltime is None and args.job_id is None:
    logger.warning('No options: -w %s or -j %s, setting %s to %s',
                   style.emph('walltime'), style.emph('job_id'),
                   style.emph('walltime'), style.emph('1:00:00'))
    args.walltime = '1:00:00'

logger.info('Options\n' + '\n'.join([style.emph(option.ljust(20)) +
            '= ' + str(vars(args)[option]).ljust(10)
            for option in sorted(vars(args).keys()) if vars(args)[option]
            or option == 'packages_management']))

# DEFINING VMS AND ELEMENTS
if args.infile:
    # parse the XML file given in arguments
    logger.info('Using %s for the topology', style.emph(args.infile))
    vm5k = parse(args.infile).getroot()
    vms = []
    for vm in vm5k.findall('.//vm'):
        vms.append(define_vms([vm.get('id')], template=vm)[0])
    if logger.getEffectiveLevel() <= 10:
        dump(vm5k)
    elements = {cluster.get('id'): len(cluster.findall('./host'))
                for cluster in vm5k.findall('.//cluster')}
else:
    logger.info('Defining VMs from template %s', style.emph(args.vm_template))
    template = fromstring(args.vm_template)
    if 'backing_file' not in template.attrib:
        template.attrib['backing_file'] = args.vm_backing_file
    vms = define_vms(['vm-' + str(i + 1) for i in range(args.n_vm)],
                     template=template)
    elements = {}
    for element in args.resources.split(','):
        if ':' in element:
            element_uid, n_nodes = element.split(':')
        else:
            element_uid, n_nodes = element, 0
        elements[element_uid] = int(n_nodes)

execution_time['1-INIT'] = timer.elapsed()

##############################################################################
# MANAGING RESERVATION
##############################################################################
timer = Timer()
print_step('Reservation')
frontend = None
if args.job_id is None:
    show_resources(elements, 'Resources wanted')
    logger.info('Finding a slot for your reservation')
    if args.kavlan:
        kavlan = True
        subnet = False
        elements['kavlan'] = 1
    else:
        kavlan = False
        subnet = True
        subnets = 'slash_22=' + str(int(ceil(len(vms) / 1024.)))

    if args.blacklisted is not None:
        blacklisted = list(set(blacklisted + args.blacklisted.split(',')))

    logger.debug('Blacklisted elements : ' + pformat(blacklisted))
    planning = get_planning(elements, vlan=kavlan, subnet=subnet)
    slots = compute_slots(planning, walltime=args.walltime,
                          excluded_elements=blacklisted)

    # Test if we need a free slot or a vms slot
    if len([element for element, n_nodes in elements.iteritems()
            if n_nodes > 0 and element != 'kavlan']) > 0:
        slot = find_free_slot(slots, elements)
        startdate = slot[0]
        resources = distribute_hosts(slot[2], elements,
                                     excluded_elements=blacklisted)
    else:
        startdate, resources = get_vms_slot(vms, elements, slots,
                                            excluded_elements=blacklisted)

    if startdate is None:
        logger.error('Unable to find a slot, exiting')
        exit()

    show_resources(resources)

    jobs_specs = get_jobs_specs(resources, name='VM5K',
                                excluded_elements=blacklisted)

    if not kavlan:
        for OarSubmission, _ in jobs_specs:
            OarSubmission.resources = subnets + '+' + OarSubmission.resources

    logger.debug('Jobs specifications %s', pformat(jobs_specs))
    if len(jobs_specs) > 1:
        job_id, _ = oargridsub(jobs_specs, walltime=args.walltime,
                               additional_options="-t deploy",
                               reservation_date=startdate)
        if job_id is None:
            cmd = get_oargridsub_commandline(jobs_specs,
                                             walltime=args.walltime,
                                             additional_options="-t deploy",
                                             reservation_date=startdate)
    else:
        sub, frontend = jobs_specs[0]
        sub.walltime = args.walltime
        sub.additional_options = "-t deploy"
        sub.reservation_date = startdate
        jobs = oarsub([(sub, frontend)])
        job_id = jobs[0][0]
        if len(jobs) == 0:
            cmd = get_oarsub_commandline((sub, frontend))
            exit()

    logger.info('Job %s will start at %s', style.emph(job_id),
                style.log_header(format_date(startdate)))
else:
    logger.info('Using an existing job: %s', style.emph(args.job_id))
    if ':' in args.job_id:
        jobs = []
        for oar_job in args.job_id.split(','):
            frontend, job_id = oar_job.split(':')
            jobs.append((job_id, frontend))
    else:
        frontend, job_id = None, args.job_id

execution_time['2-RESERVATION'] = timer.elapsed()


##############################################################################
# RETRIEVING RESOURCES
##############################################################################
timer = Timer()
print_step('Ressources')
if frontend is None:
    vm5k_resources = get_oargrid_job_vm5k_resources(job_id)
else:
    vm5k_resources = get_oar_job_vm5k_resources(jobs)
logger.debug(vm5k_resources)

execution_time['3-RESOURCES'] = timer.elapsed()


##############################################################################
# INSTALLING THE HOSTS
##############################################################################
timer = Timer()
vm5k = vm5k_deployment(infile=args.infile,
                       resources=vm5k_resources,
                       vms=vms,
                       distribution=args.vm_distribution,
                       env_name=args.env_name,
                       env_file=args.env_file,
                       outdir=args.outdir)

print_step('Deploying the hosts')
if args.nodeploy:
    vm5k._launch_kadeploy(max_tries=0)
    logger.info('Skipping packages management and libvirt configuration')
else:
    vm5k.hosts_deployment(max_tries=deployment_tries,
                          check_deploy=not args.forcedeploy)
    if args.packages_management:
        print_step('Managing packages')
        vm5k.packages_management(upgrade=True,
                                 other_packages=args.other_packages,
                                 launch_disk_copy=True,
                                 apt_cacher=args.aptcacher)
    else:
        logger.info('Not managing packages for libvirt')
        other_packages = ' '.join(args.other_packages.split(',')) \
            + 'netcat-traditional' \
            if args.other_packages else 'netcat-traditional'
        vm5k._other_packages(other_packages)
        vm5k._start_disk_copy()

    print_step('Configuring libvirt')
    vm5k.configure_libvirt()

# Saving the list of hosts in outdir
f = open(args.outdir + '/hosts.list', 'w')
for host in vm5k.hosts:
    f.write(host + '\n')
f.close()


print_step('Configuring service node')
vm5k.configure_service_node()

execution_time['4-HOSTS'] = timer.elapsed()


##############################################################################
# DEPLOYING THE VIRTUAL MACHINES
##############################################################################
timer = Timer()
print_step('Deploy virtual machines')
if args.infile is None:
    logger.info('Maximum number of VMs %s',
                get_max_vms(vm5k.hosts,
                            int(fromstring(args.vm_template).get('mem'))))

# gr = g5k_graph()
# for host in vm5k.hosts:
#     gr.add_host(host)
# for vm in vm5k.vms:
#     gr.add_node(vm['id'], {'kind': 'vm'})
#     gr.add_edge(vm['id'], vm['host'])
# plot = treemap(gr)
# plot.savefig('test.png')

f = open(args.outdir + '/vms.list', 'w')
for vm in vm5k.vms:
    f.write(vm['ip'] + '\t' + vm['id'] + '\n')
f.close()
vm5k.get_state(name='initial_topo')
vm5k.deploy_vms(clean_disks=args.vm_clean_disks,
                disk_location=args.vm_disk_location,
                apt_cacher=args.aptcacher)
vm5k.get_state(name='final_topo', plot=args.plot)

execution_time['5-VMS'] = timer.elapsed()

total = sum(value for value in execution_time.itervalues())
print_step('VM5K Execution terminated in ' + str(format_duration(total)))
log = ''

for step in sorted(execution_time.keys()):
    log += '\n' + style.emph(step + ':').ljust(10) + \
        format_duration(execution_time[step])

logger.info('Details: \nFiles saved in %s %s', style.emph(args.outdir), log)


##############################################################################
# LAUNCHING PROGRAM
##############################################################################
if args.program is not None:
    print_step('Lauching program')
    logger.info(args.program)
    if args.program in os.listdir('.'):
        args.program = './' + args.program
    prog = Process(args.program)
    prog.shell = True
    prog.stdout_handlers.append(sys.stdout)
    prog.stderr_handlers.append(sys.stderr)
    prog.run()
