#!/usr/bin/env python

"""
Jenkins runs this script to perform a test of an acme
test suite. Essentially, a wrapper around create_test and
wait_for_tests that handles cleanup of old test results and
ensures that the batch system is left in a clean state.
"""

import acme_util
acme_util.check_minimum_python_version(2, 7)
import wait_for_tests

import argparse, sys, os, shutil, glob, doctest, signal

from acme_util import expect, warning, verbose_print

SENTINEL_FILE = "ONGOING_TEST"

# Don't know if this belongs here longterm
MACHINES_THAT_MAINTAIN_BASELINES = ("redsky", "melvin", "skybridge")

###############################################################################
def parse_command_line(args, description):
###############################################################################
    parser = argparse.ArgumentParser(
usage="""\n%s [-g] [-d] [--verbose]
OR
%s --help
OR
%s --test

\033[1mEXAMPLES:\033[0m
    \033[1;32m# Run the tests and compare baselines \033[0m
    > %s
    \033[1;32m# Run the tests, compare baselines, and update dashboard \033[0m
    > %s -d
    \033[1;32m# Run the tests, generating a full set of baselines (useful for first run on a machine) \033[0m
    > %s -g
""" % ((os.path.basename(args[0]), ) * 6),

description=description,

formatter_class=argparse.ArgumentDefaultsHelpFormatter
)

    machine = acme_util.probe_machine_name()
    default_test_suite = acme_util.get_machine_info("TESTS")
    default_maintain_baselines = machine in MACHINES_THAT_MAINTAIN_BASELINES

    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Print extra information")

    parser.add_argument("-g", "--generate-baselines", action="store_true",
                        help="Generate baselines")

    parser.add_argument("--baseline-compare", action="store", choices=("yes", "no"), default=("yes" if default_maintain_baselines else "no"),
                        help="Do baseline comparisons")

    parser.add_argument("-d", "--submit-to-cdash", action="store_true",
                        help="Send results to CDash")

    parser.add_argument("-c", "--cdash-build-name",
                        help="Build name to use for CDash submission. Default will be <TEST_SUITE>_<BRANCH>_<COMPILER>")

    parser.add_argument("-p", "--cdash-project", default=wait_for_tests.ACME_MAIN_CDASH,
                        help="The name of the CDash project where results should be uploaded")

    parser.add_argument("-b", "--baseline-name", default=acme_util.get_current_branch(repo=acme_util.get_acme_root()),
                        help="Baseline name for baselines to use. Also impacts dashboard job name. Useful for testing a branch other than next or master")

    parser.add_argument("-t", "--test-suite", default=default_test_suite,
                        help="Override default acme test suite that will be run")

    parser.add_argument("--cdash-build-group", default=wait_for_tests.CDASH_DEFAULT_BUILD_GROUP,
                        help="The build group to be used to display results on the CDash dashboard.")

    args = parser.parse_args(args[1:])

    acme_util.set_verbosity(args.verbose)

    expect(not (args.submit_to_cdash and args.generate_baselines),
           "Does not make sense to use -g and -d together")
    expect(not (args.cdash_build_name is not None and not args.submit_to_cdash),
           "Does not make sense to use --cdash-build-name without -d")
    expect(not (args.cdash_project is not wait_for_tests.ACME_MAIN_CDASH and not args.submit_to_cdash),
           "Does not make sense to use -p without -d")

    return args.generate_baselines, args.submit_to_cdash, args.baseline_name, args.cdash_build_name, args.cdash_project, args.test_suite, args.cdash_build_group, args.baseline_compare

###############################################################################
def cleanup_queue(set_of_jobs_we_created):
###############################################################################
    """
    Delete all jobs left in the queue
    """
    current_jobs = set(acme_util.get_my_queued_jobs())
    jobs_to_delete = set_of_jobs_we_created & current_jobs

    if (jobs_to_delete):
        warning("Found leftover batch jobs that need to be deleted: %s" % ", ".join(jobs_to_delete))
        stat = acme_util.delete_jobs(jobs_to_delete)[0]
        if (stat != 0):
            warning("FAILED to clean up leftover jobs!")

###############################################################################
def jenkins_generic_job(generate_baselines, submit_to_cdash, baseline_name,
                        arg_cdash_build_name, cdash_project,
                        arg_test_suite,
                        cdash_build_group, baseline_compare):
###############################################################################
    """
    Return True if all tests passed
    """
    use_batch = acme_util.does_machine_have_batch()
    compilers, test_suite, scratch_root, proxy = \
        acme_util.get_machine_info(["COMPILERS", "TESTS", "CESMSCRATCHROOT", "PROXY"])
    compiler = compilers[0]
    test_suite = test_suite if arg_test_suite is None else arg_test_suite
    test_root = os.path.join(scratch_root, "jenkins")

    if (use_batch):
        batch_system = acme_util.get_batch_system()
        expect(batch_system is not None, "Bad XML. Batch machine has no batch_system configuration.")

    #
    # Env changes
    #

    if (submit_to_cdash and proxy is not None):
        os.environ["http_proxy"] = proxy

    #
    # Update submodules (Jenkins is struggling with this at the moment)
    #

    acme_util.run_cmd("git submodule update --init", from_dir=acme_util.get_acme_root())

    #
    # Cleanup previous test leftovers. Code beyond here assumes that
    # we have the scratch_root and test_root to ourselves. No other ACME testing
    # should be happening for the current user. A sentinel file helps
    # enforce this.
    #
    # Very tiny race window here, not going to sweat it
    #

    sentinel_path = os.path.join(scratch_root, SENTINEL_FILE)
    expect(not os.path.isfile(sentinel_path),
           "Tests were already in progress, cannot start more!")

    if (not os.path.isdir(scratch_root)):
        os.makedirs(scratch_root)

    open(sentinel_path, 'w').close()

    try:

        # Important, need to set up signal handlers before we officially
        # kick off tests. We don't want this process getting killed outright
        # since it's critical that the cleanup in the finally block gets run
        wait_for_tests.set_up_signal_handlers()

        #
        # Clean up leftovers from previous run of jenkins_generic_job
        #

        # Remove the old CTest XML
        if (os.path.isdir("Testing")):
            shutil.rmtree("Testing")

        # Remove the old build/run dirs
        test_id_root = "jenkins_%s" % baseline_name
        for old_dir in glob.glob("%s/*%s*" % (scratch_root, test_id_root)):
            shutil.rmtree(old_dir)

        # Remove the old cases
        for old_file in glob.glob("%s/*%s*" % (test_root, test_id_root)):
            if (os.path.isdir(old_file)):
                shutil.rmtree(old_file)
            else:
                os.remove(old_file)

        #
        # Make note of things already in the queue so we know not to delete
        # them if we timeout
        #

        if (use_batch):
            preexisting_queued_jobs = acme_util.get_my_queued_jobs()

        #
        # Set up create_test command and run it
        #

        baseline_args = ""
        if (generate_baselines):
            baseline_args = "-g -b %s" % baseline_name
        elif (baseline_compare == "yes"):
            baseline_args = "-c -b %s" % baseline_name

        test_id = "%s_%s" % (test_id_root, acme_util.get_utc_timestamp())
        create_test_cmd = "./create_test %s --test-root %s -t %s %s" % \
                          (test_suite, test_root, test_id, baseline_args)

        if (not wait_for_tests.SIGNAL_RECEIVED):
            create_test_stat = acme_util.run_cmd(create_test_cmd, from_dir=acme_util.get_acme_scripts_root(),
                                                 verbose=True, arg_stdout=None, arg_stderr=None, ok_to_fail=True)[0]
            # Create_test should have either passed, detected failing tests, or timed out
            expect(create_test_stat in [0, acme_util.TESTS_FAILED_ERR_CODE, -signal.SIGTERM],
                   "Create_test script FAILED with error code '%d'!" % create_test_stat)

        if (use_batch):
            # This is not fullproof. Any jobs that happened to be
            # submitted by this user while create_test was running will be
            # potentially deleted. This is still a big improvement over the
            # previous implementation which just assumed all queued jobs for this
            # user came from create_test.
            # TODO: change this to probe test_root for jobs ids
            our_jobs = set(acme_util.get_my_queued_jobs()) - set(preexisting_queued_jobs)

        #
        # Wait for tests
        #

        if (submit_to_cdash):
            cdash_build_name = "_".join([test_suite, baseline_name, compiler]) if arg_cdash_build_name is None else arg_cdash_build_name
        else:
            cdash_build_name = None

        tests_passed = wait_for_tests.wait_for_tests(glob.glob("%s/*%s*/TestStatus" % (test_root, test_id)),
                                                     not use_batch, # wait if using queue
                                                     False, # don't check throughput
                                                     False, # don't check memory
                                                     False, # don't ignore namelist diffs
                                                     cdash_build_name,
                                                     cdash_project,
                                                     cdash_build_group)
        if (not tests_passed and use_batch and wait_for_tests.SIGNAL_RECEIVED):
            # Cleanup
            cleanup_queue(our_jobs)

        return tests_passed
    finally:
        expect(os.path.isfile(sentinel_path), "Missing sentinel file")
        os.remove(sentinel_path)

###############################################################################
def _main_func(description):
###############################################################################
    if ("--test" in sys.argv):
        test_results = doctest.testmod(verbose=True)
        sys.exit(1 if test_results.failed > 0 else 0)

    acme_util.stop_buffering_output()

    generate_baselines, submit_to_cdash, cdash_build_name, cdash_project, baseline_branch, test_suite, cdash_build_group, no_baseline_compare = \
        parse_command_line(sys.argv, description)

    sys.exit(0 if jenkins_generic_job(generate_baselines, submit_to_cdash, cdash_build_name, cdash_project, baseline_branch, test_suite, cdash_build_group, no_baseline_compare)
             else acme_util.TESTS_FAILED_ERR_CODE)

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

if (__name__ == "__main__"):
    _main_func(__doc__)
