#!/usr/bin/env python

#       Licensed to the Apache Software Foundation (ASF) under one
#       or more contributor license agreements.  See the NOTICE file
#       distributed with this work for additional information
#       regarding copyright ownership.  The ASF licenses this file
#       to you under the Apache License, Version 2.0 (the
#       "License"); you may not use this file except in compliance
#       with the License.  You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#       Unless required by applicable law or agreed to in writing,
#       software distributed under the License is distributed on an
#       "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#       KIND, either express or implied.  See the License for the
#       specific language governing permissions and limitations
#       under the License.

import argparse
from copy import copy
from glob import glob
import multiprocessing
from multiprocessing.pool import ThreadPool
import subprocess
import sys
import threading
import textwrap
import os

CPUS = multiprocessing.cpu_count()
CONCURRENT_SUITES = (CPUS // 4) or CPUS
CONCURRENT_TESTS = CPUS // CONCURRENT_SUITES
PROC_TIMEOUT = 360

ALT_PKG_PATHS = {
    'Allura': 'allura/tests/',
}

NOT_MULTIPROC_SAFE = [
    'ForgeGit',
    'ForgeSVN',
]


def run_one(cmd, **popen_kwargs):
    cmd_to_show = u'{} in {}'.format(cmd, popen_kwargs.get('cwd', '.'))
    print u'{} running {}'.format(threading.current_thread(), cmd_to_show)
    sys.stdout.flush()

    all_popen_kwargs = dict(shell=True, stderr=subprocess.STDOUT,
                            stdout=subprocess.PIPE,
                            bufsize=1,  # 1 == line-buffered
                            close_fds='posix' in sys.builtin_module_names)
    all_popen_kwargs.update(popen_kwargs)
    proc = subprocess.Popen(cmd, **all_popen_kwargs)
    while proc.poll() is None:
        line = proc.stdout.readline()
        sys.stdout.write(line)
        sys.stdout.flush()
    # wait for completion and get remainder of output
    out_remainder, _ = proc.communicate()
    sys.stdout.write(out_remainder)
    sys.stdout.flush()
    print u'finished {}'.format(cmd_to_show)
    sys.stdout.flush()
    return proc


def run_many(cmds, processes=None):
    """
    cmds: list of shell commands, or list of (shell cmds, popen_kwargs)
    processes: number of processes, or None for # of CPU cores
    """
    thread_pool = ThreadPool(processes=processes)

    async_results = []
    for cmd_kwds in cmds:
        if type(cmd_kwds) == ():
            cmd = cmd_kwds
            kwds = {}
        else:
            cmd = cmd_kwds[0]
            kwds = cmd_kwds[1]
        result = thread_pool.apply_async(run_one, args=(cmd,), kwds=kwds)
        async_results.append(result)

    thread_pool.close()
    thread_pool.join()

    procs = [async_result.get() for async_result in async_results]
    return [p.returncode for p in procs]


def get_packages():
    packages = sorted([p.split('/')[0] for p in glob("*/setup.py")])

    # make it first, to catch syntax errors
    packages.remove('AlluraTest')
    packages.insert(0, 'AlluraTest')
    return packages


def check_packages(packages):
    for pkg in packages:
        try:
            __import__(pkg.lower())
        except ImportError:
            print "Not running tests for {}, since it isn't set up".format(pkg)
        else:
            yield pkg


def run_tests_in_parallel(options, nosetests_args):
    # coverage and multiproc plugins not compatible
    use_multiproc = '--with-coverage' not in nosetests_args

    def get_pkg_path(pkg):
        return ALT_PKG_PATHS.get(pkg, '')

    def get_multiproc_args(pkg):
        if not use_multiproc or options.concurrent_tests == 1:
            return ''
        return ('--processes={procs_per_suite} --process-timeout={proc_timeout}'.format(
            procs_per_suite=options.concurrent_tests,
            proc_timeout=PROC_TIMEOUT)
            if pkg not in NOT_MULTIPROC_SAFE else '')

    def get_concurrent_suites():
        if use_multiproc or '-n' in sys.argv:
            return options.concurrent_suites
        return CPUS

    cmds = []
    env = dict(os.environ,
               NOSE_IGNORE_CONFIG_FILES='1')
    for package in check_packages(options.packages):
        cover_package = package.lower()
        our_nosetests_args = copy(nosetests_args)
        our_nosetests_args.append('--cover-package={}'.format(cover_package))
        cmd = "nosetests {pkg_path} {nosetests_args} {multiproc_args}".format(
            pkg_path=get_pkg_path(package),
            nosetests_args=' '.join(our_nosetests_args),
            multiproc_args=get_multiproc_args(package),
        )
        cmds.append((cmd, dict(cwd=package, env=env)))
    return run_many(cmds, processes=get_concurrent_suites())


def parse_args():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent('''
                                        All additional arguments are passed along to nosetests
                                          (e.g. -v --with-coverage)
                                        Note: --cover-package will be set automatically to the appropriate value'''))
    parser.add_argument('-n', help='Number of test suites to run concurrently in separate '
                                   'processes. Default: # CPUs / 4',
                        dest='concurrent_suites', type=int, default=CONCURRENT_SUITES)
    parser.add_argument('-m', help='Number of tests to run concurrently in separate '
                                   'processes, per suite. Default: # CPUs / # concurrent suites',
                        dest='concurrent_tests', type=int, default=CONCURRENT_TESTS)
    parser.add_argument(
        '-p', help='List of packages to run tests on. Default: all',
        dest='packages', choices=get_packages(), default=get_packages(),
        nargs='+')
    return parser.parse_known_args()


if __name__ == "__main__":
    ret_codes = run_tests_in_parallel(*parse_args())
    sys.exit(any(ret_codes))