#!/usr/bin/python
import argparse
import tempfile
from benchpress.press import expand_path, handle_result_file

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Runs a benchmark suite and stores the results in a json-file.')
    parser.add_argument(
        'bohrium_src',
        help='Path to the Bohrium source-code.',
        type=lambda x: expand_path(parser, x)
    )
    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(
        '--no-perf',
        action="store_true",
        help="Disable the use of the perf measuring tool."
    )
    parser.add_argument(
        '--no-time',
        action="store_true",
        help="Disable 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. '
    )
    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.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)
    try:
        try:
            if args.output is None:
                raise IOError
            result_file = open(args.output, 'r+')
        except IOError:
            if args.output is None:
                result_file = tempfile.NamedTemporaryFile(delete=False,
                                             prefix=file_name_prefix,
                                             suffix='.json')
            else:
                result_file = open(args.output, 'w+')

            #Populate 'result_file' with benchmark jobs/runs
            gen_jobs(result_file, os.getenv('HOME')+os.sep+'.bohrium'+os.sep+'config.ini', 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:
        result_file.close()
