#!/usr/bin/python
import argparse
import tempfile
import os
from benchpress.press import *

class _C:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'

DUMMY_BH_CONFIG = """
[bridge]
type = bridge

[node]
type = vem

[cpu]
type = ve
"""

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Runs a benchmark suite and stores the results in a json-file.')
    parser.add_argument(
        'repos_root',
        help='Path to git-repository, revision will be included in result. Use "NOREPOS" if you do not want it.',
        type=lambda x: expand_path(parser, x) if not x == "NOREPOS" else None
    )
    parser.add_argument(
        'suite_file',
        type=argparse.FileType('r'),
        help="Path to the benchmark suite file."
    )
    parser.add_argument(
        '--output',
        type=str,
        metavar='RESULT_FILE',
        help='Path to the JSON file where the benchmark results will be written. '
             'If the file exist, the benchmark will resume.'
    )
    parser.add_argument(
        '--runs',
        default=3,
        type=int,
        help="How many times should each benchmark run."
    )
    parser.add_argument(
        '--with-perf',
        action="store_true",
        help="Enable the use of the perf measuring tool."
    )
    parser.add_argument(
        '--with-time',
        action="store_true",
        help="Enable the use of the '/usr/bin/time -v' measuring tool."
    )
    parser.add_argument(
        '--save-data',
        action="store_true",
        help="Save data output from benchmarks in RESULT_FILE. "\
             "All benchmarks must support the --outputfn argument."
    )
    parser.add_argument(
        '--pre-clean',
        action="store_true",
        help="Clean caches such as the fuse or the kernel cache before execution."
    )
    parser.add_argument(
        '--restart',
        action="store_true",
        help="Restart execution or submission of failed jobs."
    )
    parser.add_argument(
        '--publish-cmd',
        type=str,
        metavar='COMMAND',
        help='The publish command to use before exiting (use together with --wait). '\
             'NB: $OUT is replaced with the name of the output JSON file. '
    )
    parser.add_argument(
        '--bh-config-file',
        help="Override the path to the bohrium config.ini file, which is normally read from ~/.bohrium/config.ini."
    )
    parser.add_argument(
        '--resume-only',
        action="store_true",
        help="Ensures that no new jobs will be generated, the script will always resume a previous run or fail."
    )
    slurm_grp = parser.add_argument_group('SLURM Queuing System')
    slurm_grp.add_argument(
        '--slurm',
        action="store_true",
        help="Use the SLURM queuing system. This overwrite the default value "\
	     "specified in the suite ('use_slurm_default')"
    )
    slurm_grp.add_argument(
        '--no-slurm',
        action="store_true",
        help="Do not use the SLURM queuing system. This overwrite the default value "\
	     "specified in the suite ('use_slurm_default')"
    )
    slurm_grp.add_argument(
        '--partition',
        type=str,
        help="Submit to a specific SLURM partition."
    )
    slurm_grp.add_argument(
        '--multi-jobs',
        action="store_true",
        help="Submit 'runs' SLURM jobs instead of one job with 'runs' number of runs."
    )
    slurm_grp.add_argument(
        '--wait',
        action="store_true",
        help="Wait for all SLURM jobs to finished before returning."
    )
    args = parser.parse_args()

    if args.with_perf:
        args.no_perf = False
    if args.with_time:
        args.no_timer = False

    if args.publish_cmd and not args.wait:
        parser.error( '--publish-cmd must be used together with --wait' )

    file_name_prefix = 'benchmark-%s-' % os.path.basename(args.suite_file.name)
    temp_bh_config = None
    result_file = None

    try:
        try:
            if args.output is None:
                raise IOError
            result_file = open(args.output, 'r+')
        except IOError:
            if args.resume_only:
                raise Exception('Failed to open output file in resume-only mode')

            if args.output is None:
                result_file = tempfile.NamedTemporaryFile(delete=False,
                                             prefix=file_name_prefix,
                                             suffix='.json')
            else:
                result_file = open(args.output, 'w+')

            bh_config_file = args.bh_config_file

            if bh_config_file == None:

                guesses = [
                          os.getenv('BH_CONFIG'),
                          os.path.join(os.getenv('HOME'),     '.bohrium', 'config.ini'),
                    '/' + os.path.join('etc',                 'bohrium', 'config.ini'),
                    '/' + os.path.join('usr', 'etc',          'bohrium', 'config.ini'),
                    '/' + os.path.join('usr', 'local', 'etc', 'bohrium', 'config.ini'),
                    '/' + os.path.join('opt', 'etc',          'bohrium', 'config.ini')
                ]

                for guess in guesses:
                    if guess != None and os.path.isfile(guess):
                        bh_config_file = guess
                        print _C.OKBLUE, "Using Bohrium config file",bh_config_file,_C.ENDC
                        break

            if bh_config_file == None:
                handle, temp_bh_config = tempfile.mkstemp(prefix=file_name_prefix, suffix='.ini')
                bh_config_file = temp_bh_config
                print _C.WARNING, "Bohrium config file not found, using dummy file",bh_config_file,_C.ENDC

                with os.fdopen(handle, 'w+') as config_file:
                    config_file.write(DUMMY_BH_CONFIG)

            if not os.path.isfile(bh_config_file):
                raise IOError('File not found: ' + bh_config_file)

            #Populate 'result_file' with benchmark jobs/runs
            gen_jobs(result_file, bh_config_file, args)

        if args.wait:
            while True:#Probe for finished jobs (one second delay)
                if handle_result_file(result_file, args):
                    all_finished = True
                    break
                time.sleep(5)
        else:
            all_finished = handle_result_file(result_file, args)

        if all_finished:
            print _C.OKGREEN,"Benchmark all finished:",result_file.name,_C.ENDC
        else:
            print _C.WARNING,"Benchmark saved in",result_file.name,
            print "use resume on",result_file.name,_C.ENDC

        if args.publish_cmd:
            cmd = args.publish_cmd.replace("$OUT", result_file.name)
            print _C.WARNING,"Publishing results using '%s'"%cmd,_C.ENDC
            check_call(cmd, shell=True)

    except KeyboardInterrupt:
        print _C.WARNING,"Suspending the benchmark execution,",
        print "use resume on",result_file.name,_C.ENDC
    finally:
        if result_file != None:
            try:
                result_file.close()
            except:
                pass

        if temp_bh_config != None:
            try:
                os.unlink(temp_bh_config)
            except:
                pass

