#!/usr/bin/env python
"""
Runner Pyzza Launcher Manager

RPlauncher

Runner Pyzza main script 
"""
from RunnerPyzza import __version__
from RunnerPyzza.ClientCommon.PyzzaTalk import OrderPyzza, OvenPyzza, CheckPyzza, \
    EatPyzza, CleanPyzza
from RunnerPyzza.LauncherManager.XMLHandler import MachinesSetup, ScriptChain
from RunnerPyzza.Common.ColorLog import ColorFormatter
import argparse
import getpass
import logging
import sys
import time

__author__ = "Emilio Potenza"
__credits__ = ["Marco Galardini"]

################################################################################
# Log

logger = logging.getLogger('RunnerPyzza')

################################################################################
# Methods

def auto(options):
    logger.info('Automatic pyzza!')
    
    if not init(options):
        return
    
    time.sleep(1)
    
    if not start(options):
        return
    
    time.sleep(1)
    
    while True:
        if not status(options):
            return

        if options.status == 'ready':
            break
        
        time.sleep(5)
            
    if not results(options):
        return
    
    if not clean(options):
        return

def init(options):
    logger.info("Reading inputs")
    f=open(options.scriptChain)
    h = ScriptChain(''.join(f.readlines()))
    
    # No machines can be added if we have an scp job
    if options.machines != '' and options.local:
        logger.warning('No machines should be provided for a local job')
    
    if options.machines != '':
        logger.info("Reading machines XML...")
        m = MachinesSetup(options.machines)
        machs = m.getMachines()
    else:
        logger.info("No machines provided...")
        machs = []

    if options.local and options.password:
        password = getpass.getpass('%s password for %s: '%
                                   (options.user, options.host))
    else:
        password = ''

    logger.info("Open cominication with daemon...")
    try:
        order = OrderPyzza(options.host, options.port,
                machines = machs, programs = h.getPrograms(),
                tag = options.tag, local = options.local, localdir = options.indir,
                user = options.user, password = password)
    except Exception, e:
        logger.warning('Could not order the pyzza (%s)'%e)
        return False
    
    if not order.launchOrder():
        logger.warning('Pyzza not ordered!')
        return False
    else:
        logger.info('Pyzza ordered with id %s'%order.jobID)
    
    setattr(options, 'jobID', order.jobID)
    
    return True

def start(options):
    logger.info('Let\'s put the pyzza in the oven!')
    ovenizer = OvenPyzza(options.host, options.port, options.jobID)
    if not ovenizer.putInTheOven():
        logger.warning('The oven is cold! Could not cook %s'%options.jobID)
        return False
    
    logger.info('Pyzza with id %s is in the oven!'%options.jobID)
    
    return True

def status(options):
    logger.info('Let\'s check if the pyzza is ready!')
    checker = CheckPyzza(options.host, options.port, options.jobID)
    if not checker.checkTheOven():
        logger.warning('Could not find the pyzza! %s'%options.jobID)
        return False
    else:
        if checker.isReady():
            logger.info('The Pyzza with id %s is ready!'%options.jobID)
            setattr(options,'status','ready')
        elif checker.isCooking():
            logger.warning('The Pyzza with id %s is still in the oven!'%options.jobID)
            logger.warning('%d slices cooked!'%(int(checker.getLastSlice()[1])))
            setattr(options,'status','running')
        elif checker.isWaiting():
            logger.warning('The oven is still cold! Waiting to cook %s'%options.jobID)
            setattr(options,'status','waiting')
        elif checker.isBurned():
            logger.error('The Pyzza with id %s is burned!'%options.jobID)
            logger.error('Problematic slices:')
            for err in checker.inspectErrors():
                logger.error('Slice %d, Status %s, Ingredients %s, Return code %d'%(int(err[1]),err[0],err[2],int(err[3])))
            return False
    
    return True

def results(options):
    logger.info('Let me eat my pyzza!')
    
    if options.local and options.password:
        password = getpass.getpass('%s password for %s: '%
                                   (options.user, options.host))
    else:
        password = ''
        
    eater = EatPyzza(options.host, options.port, options.jobID,
                     local = options.local, user = options.user,
                     password = password)
    if not eater.eatThePyzza():
        logger.error('Could not eat the pyzza! %s'%options.jobID)
        return False
    else:
        # Open the output file
        if options.out == '':
            fname = options.jobID + '.tsv'
        else:
            fname = options.out
        logger.info('Saving slices outputs to %s'%fname)
        fout = open(fname,'w')
        fout.write( '\t'.join( ['# Slice', 'Name', 'Machine', 'Exit code',
                                'stdout', 'stderr'] ) )
        fout.write('\n')
        
        for pyzzaslice in eater.getSlices():
            logger.info('Eating slice %d'%pyzzaslice)
            for program in eater.eatSlice(pyzzaslice):
                logger.info('Results for program %s'%program.name)
                logger.info('Machine %s'%program.getHost())
                logger.info('Exit status %d'%program.getExit())
                for line in program.getStdOut().splitlines():
                    fout.write( '\t'.join( [str(pyzzaslice), program.name,
                                    program.getHost(), str(program.getExit()),
                                    line, ''] ) )
                    fout.write('\n')
                    logger.info('\033[1;32m[%s]\033[0m'%line)
                for line in program.getStdErr().splitlines():
                    fout.write( '\t'.join( [str(pyzzaslice), program.name,
                                    program.getHost(), str(program.getExit()),
                                    '', line] ) )
                    fout.write('\n')
                    logger.info('\033[1;31m[%s]\033[0m'%line)
        
        fout.close()
        
    return True

def clean(options):
    logger.info('Let\'s clean the table and go home!')
    cleaner = CleanPyzza(options.host, options.port, options.jobID)
    if not cleaner.cleanAndPay():
        logger.error('Could not clean the table and pay the pyzza! %s'%options.jobID)
        return False
    
    return True

################################################################################
# Read options

def getOptions():
    # create the top-level parser
    description = "RPlauncher.py %s RunnerPyzza main script"%__version__
    parser = argparse.ArgumentParser(description = description)
    parser.add_argument('-d', '--host', action='store',
                        default='localhost',
                        help='Main server hostname')
    parser.add_argument('-p', '--port', action='store',
                        default=4123,
                        help='Main server port')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='Suppress verbosity')
    parser.add_argument('-g', '--debug', action='store_true',
                        help='Add debug messages')
    subparsers = parser.add_subparsers()

    parser_auto = subparsers.add_parser('auto', help='Eat a pyzza')
    parser_auto.add_argument('scriptChain', action="store",
                            help='ScriptChain file')
    parser_auto.add_argument('-m', '--machines', action="store",
                             default='',
                            help='Machines file')
    parser_auto.add_argument('-t', '--tag', action="store",
                             default='Margherita',
                            help='Pyzza tag')
    parser_auto.add_argument('-o', '--out', action="store",
            default='',
            dest='out',
            help='Results file name [Default: job ID]')
    parser_auto.add_argument('-l', '--local', action="store_true",
            default=False,
            help='Make a copy of the local input files on the main server (if NFS is not available)')
    parser_auto.add_argument('-i', '--input-dir', action="store",
            default='',
            dest='indir',
            help='Input directory to be transferred on the main server (only if -l)')
    parser_auto.add_argument('-u', '--user', action="store",
            default='runnerpyzza',
            dest='user',
            help='User for the input-directory transfer [default: runnerpyzza] (only if -l)')
    parser_auto.add_argument('-p', '--password', action="store_true",
            default=False,
            dest='password',
            help='Use a password for the input-directory transfer (only if -l)')
    parser_auto.set_defaults(func=auto)

    parser_init = subparsers.add_parser('init', help='Order a pyzza')
    parser_init.add_argument('scriptChain', action="store",
                            help='ScriptChain file')
    parser_init.add_argument('-m', '--machines', action="store",
                             default='',
                            help='Machines file')
    parser_init.add_argument('-t', '--tag', action="store",
                             default='Margherita',
                            help='Pyzza tag')
    parser_init.add_argument('-l', '--local', action="store_true",
            default=False,
            help='Make a copy of the local input files on the main server (if NFS is not available)')
    parser_init.add_argument('-i', '--input-dir', action="store",
            default='',
            dest='indir',
            help='Input directory to be transferred on the main server (only if -l)')
    parser_init.add_argument('-u', '--user', action="store",
            default='runnerpyzza',
            dest='user',
            help='User for the input-directory transfer [default: runnerpyzza] (only if -l)')
    parser_init.add_argument('-p', '--password', action="store_true",
            default=False,
            dest='password',
            help='Use a password for the input-directory transfer (only if -l)')
    parser_init.set_defaults(func=init)

    parser_start = subparsers.add_parser('start', help='Put the pyzza in the oven')
    parser_start.add_argument('jobID', action="store",
                            help='Job ID')
    parser_start.set_defaults(func=start)
    
    parser_status = subparsers.add_parser('status', help='Check the pyzza')
    parser_status.add_argument('jobID', action="store",
                            help='Job ID')
    parser_status.set_defaults(func=status)
    
    parser_results = subparsers.add_parser('results', help='Eat the pyzza')
    parser_results.add_argument('jobID', action="store",
                            help='Job ID')
    parser_results.add_argument('-o', '--out', action="store",
            default='',
            dest='out',
            help='Results file name [Default: job ID]')
    parser_results.add_argument('-l', '--local', action="store_true",
            default=False,
            help='Make a copy of the results files from the main server (if NFS is not available)')
    parser_results.add_argument('-u', '--user', action="store",
            default='runnerpyzza',
            dest='user',
            help='User for the results-directory transfer [default: runnerpyzza] (only if -l)')
    parser_results.add_argument('-p', '--password', action="store_true",
            default=False,
            dest='password',
            help='Use a password for the results-directory transfer (only if -l)')
    parser_results.set_defaults(func=results)
    
    parser_clean = subparsers.add_parser('clean', help='Clean the table')
    parser_clean.add_argument('jobID', action="store",
                            help='Job ID')
    parser_clean.set_defaults(func=clean)
    
    return parser.parse_args()

################################################################################

options = getOptions()

# Set root log level
if options.quiet:
    logger.setLevel(logging.WARN)
elif options.debug:
    logger.setLevel(logging.DEBUG)
else:
    logger.setLevel(logging.INFO)

# create console handler
ch = logging.StreamHandler()
# Set log level
if options.quiet:
    ch.setLevel(logging.WARN)
elif options.debug:
    ch.setLevel(logging.DEBUG)
else:
    ch.setLevel(logging.INFO)

# create formatter for console handler
formatter = ColorFormatter('%(asctime)s - $COLOR%(message)s$RESET','%H:%M:%S')
# add formatter to console handler
ch.setFormatter(formatter)

# add console handler to logger
logger.addHandler(ch)

# create file handler
fh = logging.FileHandler('RPlauncher.log')
# Set log level
# The file handler by default print all the levels but DEBUG
if options.debug:
    fh.setLevel(logging.DEBUG)
else:
    fh.setLevel(logging.INFO)
# create formatter
formatter = logging.Formatter('%(asctime)s - [%(levelname)s] - %(message)s',
                            '%Y-%m-%d %H:%M:%S')
fh.setFormatter(formatter)
logger.addHandler(fh)

options.func(options)

################################################################################
# Implementation

def launcherQuit(client_socket = None, quit_msg = None, exit = True):
    '''
    Close RPlauncher and socket if available
    '''
    if client_socket:
        ##Close connection and start queue
        client_socket.send(quit_msg)
        # wait more than a little bit
        time.sleep(1)
        # close socket connection without error
        client_socket.close()
        logging.info("Bye Bye!")
    if exit:
        sys.exit()







