#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
firebat.console
~~~~~~~~~~~~~~~

Helper script for Phantom load tool.
    * generate config files and input data.
    * runs test.
    * aggregate result data.
"""

import os
import sys
import datetime
import argparse
import logging
import getpass
import cPickle
import commands
import copy
import base64
from progressbar import Bar, ProgressBar, Percentage, ETA

import yaml
import simplejson as json

from firebat import __version__
from firebat.console.conf import make_conf, get_defaults
from firebat.console.stepper import parse_ammo, process_load_schema
from firebat.console.stepper import fire_duration
from firebat.console.cmd import get_running_jobs
from firebat.console.aggr import process_phout


def exit_err(msg):
    logger = logging.getLogger('firebat.console')
    if isinstance(msg, str):
        msg = [msg, ]
    for m in msg:
        logger.error(m)
    if not logger.handlers:
        sys.stderr.write(msg)
    sys.exit(1)


def validate_dict(d, req):
    ''' Check that all keys from required list present in tested dict.
    '''
    present_keys = d.keys()
    diff = [val for val in req if val not in present_keys]
    if len(diff) != 0:
        raise ValueError('You missed required options in conf: %s' % diff)


def __build_path(orig_wd, config, fire, time):
    test_path = orig_wd + '/'
    test_path += config['title']['task'] + '_'
    test_path += getpass.getuser() + '_'
    test_path += time.strftime('%Y%m%d-%H%M%S') + '/'
    test_path += fire['name']
    return test_path


def start_daemon(fire):
    ''' Build cmd string for system start-stop-daemon tool and call it.
    Args:
        fire: dict, fire from job configuration.

    Returns:
        status_d: int, command exit code.
        text_d: str, command stdout.
    '''

    exec_name = 'daemon_fire'
    status, exec_path = commands.getstatusoutput('which %s' % exec_name)
    if status != 0:
        raise ValueError('Can\'t find executable: %s' % exec_name)
    opts = {
        'owner': fire.get('owner', 'uid'),
        'total_dur': fire['total_dur'] / 1000,  # in seconds
    }
    opts_str = base64.b64encode(cPickle.dumps(opts))
    cmd = 'start-stop-daemon --start --quiet --background'
    #cmd += ' --make-pidfile %s' % pid_file
    cmd += ' --pidfile /None_existent'
    cmd += ' --chdir %s' % fire['wd']
    cmd += ' --exec %s -- %s' % (exec_path, opts_str)

    status_d, text_d = commands.getstatusoutput(cmd)
    return status_d, text_d


def get_arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', nargs=1, action='store',
                    dest='config_file', default='./fire.yaml',
                    help='path to configuration file',)

    parser.add_argument('-a', '--ammo', action='store',
                    dest='ammo_file', default=None,
                    help='path to ammo file',)

    parser.add_argument('-o', '--only-prepare-stpd', action='store',
                    dest='stpd_ammo_only', default=False,
                    help='Only generate ammo stpd and exit',)

    parser.add_argument('--list', action="store_const", const=True,
            help='List details of currently running fires.', dest='list_only')

    parser.add_argument('--phout', action="store_const", const=True,
            help='Parse phout file.', dest='phout')

    parser.add_argument('--debug', action="store_const", const=True,
            help='Whether to show debug msg in STDOUT', dest='debug')

    parser.add_argument('--version', action='version',
            version='%(prog)s ' + __version__)
    return parser


def get_logger(args):
    logger = logging.getLogger('firebat.console')
    logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    if args.debug:
        ch.setLevel(logging.DEBUG)
    else:
        ch.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s  %(message)s')
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    return logger


def main():
    defaults = get_defaults()
    args = get_arg_parser().parse_args()
    if args.list_only:
        get_running_jobs()
        sys.exit(0)

    if args.phout:
        with open('phout.txt', 'r') as phout_fh:
            process_phout(phout_fh)
        sys.exit(0)

    logger = get_logger(args)

    # load job configuration
    try:
        with open(args.config_file, 'r') as conf_fh:
            config = yaml.load(conf_fh)
    except IOError, e:
        exit_err('Could not read "%s": %s\n' % (args.config_file, e))
    except yaml.scanner.ScannerError, e:
        exit_err('Could not parse config file: %s\n%s' % (args.config_file, e))

    try:
        validate_dict(config['title'], defaults['title_required_keys'])
    except ValueError, e:
        exit_err('Error in parsing fire conf:\n%s' % e)

    now = datetime.datetime.now()
    orig_wd = os.getcwd()
    ammo_from_arg = False
    for idx, f_orig in enumerate(config['fire']):
        f = defaults['fire_conf']
        f.update(f_orig)
        try:
            validate_dict(f, defaults['fire_required_keys'])
        except ValueError, e:
            exit_err('Error in parsing fire conf:\n%s' % e)

        # get absolute ammo path before chdir
        if not ammo_from_arg:
            if args.ammo_file:
                ammo_from_arg = os.path.abspath(args.ammo_file)
                ammo_path = ammo_from_arg
            else:
                ammo_path = f['input_file']

        total_dur = fire_duration(f)
        f['total_dur'] = total_dur
        pbar_max = total_dur
        widgets = [Percentage(), Bar(), ETA(), ]
        pbar = ProgressBar(widgets=widgets, maxval=pbar_max).start()

        job_path = __build_path(orig_wd, config, f, now)
        f['wd'] = job_path
        config['fire'][idx] = copy.copy(f)
        os.makedirs(job_path)
        os.chdir(job_path)

        fire_json = json.dumps(config['fire'][idx], indent=4 * ' ')
        with open('.fire.json', 'w+') as fire_fh:
            fire_fh.write(fire_json)

        with open('phantom.conf', 'w+') as cfg_fh:
            cfg_fh.write(make_conf(f))

        logger.info('Processing fire: %s' % f['name'])
        logger.info('Ammo file: %s' % ammo_path)
        logger.info('Load schema: %s' % f['load'])
        stpd_start = datetime.datetime.now()
        with open(ammo_path, 'r') as ammo_fh,\
            open('ammo.stpd', 'w+') as stpd_fh:
            ammos = parse_ammo(ammo_fh)
            offset = f.get('offset', 0)
            for schema in f['load']:
                for tick in process_load_schema(schema, offset):
                    try:
                        m_line, chunk = ammos.next()
                    except StopIteration:
                        # chunks in ammo file run out
                        if f['loop_ammo']:
                            # need to restart requests chunks generation
                            ammo_fh.seek(0)
                            ammos = parse_ammo(ammo_fh)
                            m_line, chunk = ammos.next()
                        else:
                            stpd_fh.write('0')
                            msg = [
                                'Not enough requests in ammo file' +
                                ' to cover load schema.',
                                'File: %s' % f['input_file'],
                                'Schema: %s' % schema,
                            ]
                            logger.info(msg)
                            break
                    req_stpd = m_line % tick + chunk + '\n'
                    stpd_fh.write(req_stpd)
                    pbar.update(tick)
                offset = tick  # use last time tick as next load schema offset
            stpd_fh.write('0')  # close stpd file according format
        pbar.finish()
        stpd_stop = datetime.datetime.now()
        stpd_job_dur = stpd_stop - stpd_start
        logger.info('ammo job takes: %s\n' % stpd_job_dur)
    logger.info('stpd generation finished.')

    if args.stpd_ammo_only:  # generation only, We no need to run tests
        sys.exit(0)

    # ammo stpd and configs are ready, let's start Phantom supervisor daemon.
    for f in config['fire']:
        retcode, out = start_daemon(f)
        if retcode == 0:
            logger.info('Fire %s launched successfully.' % f['name'])
        else:
            logger.error('Fire start fails: %s.Exit code: %s' % (out, retcode))

if __name__ == '__main__':
    main()
