#!/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
from multiprocessing.pool import ThreadPool
import subprocess
import sys
import threading
import textwrap


def run_one(cmd, **popen_kwargs):
    print '{} running {} {}'.format(threading.current_thread(), cmd, popen_kwargs)
    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 'finished {} {}'.format(cmd, popen_kwargs)
    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 = [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):
    cmds = []
    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 {nosetests_args}".format(
            nosetests_args=' '.join(our_nosetests_args),
        )
        cmds.append((cmd, dict(cwd=package)))
    return run_many(cmds, processes=options.num_processes)


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 processes to use at once. Default: # CPUs',
                        dest='num_processes', type=int, default=None)
    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))