#!/usr/bin/env python
from lib.PerfKit import PerfKit
from lib.Rally import Rally
from lib.Shaker import Shaker
from lib.WorkloadBase import WorkloadBase
import argparse
import logging
import sys
import yaml
import datetime
import os
from pykwalify import core as pykwalify_core
from pykwalify import errors as pykwalify_errors

_workload_opts = ['perfkit', 'rally', 'shaker']
_config_file = 'browbeat-config.yaml'
debug_log_file = 'log/debug.log'


def _load_config(path, _logger):
    try:
        stream = open(path, 'r')
    except IOError:
        _logger.error("Configuration file {} passed is missing".format(path))
        exit(1)
    config = yaml.load(stream)
    stream.close()
    validate_yaml(config, _logger)
    return config

def validate_yaml(config, _logger):
    _logger.info("Validating the configuration file passed by the user")
    stream = open("lib/validate.yaml", 'r')
    schema = yaml.load(stream)
    check = pykwalify_core.Core(source_data=config, schema_data=schema)
    try:
        check.validate(raise_exception=True)
        _logger.info("Validation successful")
    except pykwalify_errors.SchemaError as e:
        _logger.error("Schema Validation failed")
        raise Exception('File does not conform to schema: {}'.format(e))

def _run_workload_provider(provider, config):
    _logger = logging.getLogger('browbeat')
    if provider == "perfkit":
        perfkit = PerfKit(config)
        perfkit.start_workloads()
    elif provider == "rally":
        rally = Rally(config)
        rally.start_workloads()
    elif provider == "shaker":
        shaker = Shaker(config)
        shaker.run_shaker()
    else:
        _logger.error("Unknown workload provider: {}".format(provider))


def main():
    parser = argparse.ArgumentParser(
        description="Browbeat Performance and Scale testing for Openstack")
    parser.add_argument(
        '-s',
        '--setup',
        nargs='?',
        default=_config_file,
        help='Provide Browbeat YAML configuration file. Default is ./{}'.format(_config_file))
    parser.add_argument('workloads', nargs='*', help='Browbeat workload(s). Takes a space separated'
                        ' list of workloads ({}) or \"all\"'.format(', '.join(_workload_opts)))
    parser.add_argument('--debug', action='store_true', help='Enable Debug messages')
    _cli_args = parser.parse_args()

    _logger = logging.getLogger('browbeat')
    _logger.setLevel(logging.DEBUG)
    _formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)7s - %(message)s')
    _dbg_file = logging.FileHandler(debug_log_file)
    _dbg_file.setLevel(logging.DEBUG)
    _dbg_file.setFormatter(_formatter)
    _ch = logging.StreamHandler()
    if _cli_args.debug:
        _ch.setLevel(logging.DEBUG)
    else:
        _ch.setLevel(logging.INFO)
    _ch.setFormatter(_formatter)
    _logger.addHandler(_dbg_file)
    _logger.addHandler(_ch)

    _logger.debug("CLI Args: {}".format(_cli_args))

    # Load Browbeat yaml config file:
    _config = _load_config(_cli_args.setup, _logger)

    # Default to all workloads
    if _cli_args.workloads == []:
        _cli_args.workloads.append('all')

    if len(_cli_args.workloads) == 1 and 'all' in _cli_args.workloads:
        _cli_args.workloads = _workload_opts
    invalid_wkld = [wkld for wkld in _cli_args.workloads if wkld not in _workload_opts]
    if invalid_wkld:
        _logger.error("Invalid workload(s) specified: {}".format(invalid_wkld))
        if 'all' in _cli_args.workloads:
            _logger.error("If you meant 'all' use: './browbeat.py all' or './browbeat.py'")
        exit(1)
    else:
        time_stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        _logger.info("Browbeat test suite kicked off")
        _logger.info("Running workload(s): {}".format(','.join(_cli_args.workloads)))
        for wkld_provider in _cli_args.workloads:
            if wkld_provider in _config:
                if _config[wkld_provider]['enabled']:
                    _run_workload_provider(wkld_provider, _config)
                else:
                    _logger.warning("{} is not enabled in {}".format(wkld_provider,
                                                                     _cli_args.setup))
            else:
                _logger.error("{} is missing in {}".format(wkld_provider, _cli_args.setup))
        result_dir = _config['browbeat']['results']
        WorkloadBase.print_report(result_dir, time_stamp)
        _logger.info("Saved browbeat result summary to {}".format(
            os.path.join(result_dir,time_stamp + '.' + 'report')))
        WorkloadBase.print_summary()

if __name__ == '__main__':
    sys.exit(main())