# Copyright 2020 MONAI Consortium
# Licensed 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
import inspect
import os
import sys
import time
import unittest

from monai.utils import PerfContext

results: dict = dict()


class TimeLoggingTestResult(unittest.TextTestResult):
    """Overload the default results so that we can store the results."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.timed_tests = dict()

    def startTest(self, test):  # noqa: N802
        """Start timer, print test name, do normal test."""
        self.start_time = time.time()
        name = self.getDescription(test)
        self.stream.write(f"Starting test: {name}...\n")
        super().startTest(test)

    def stopTest(self, test):  # noqa: N802
        """On test end, get time, print, store and do normal behaviour."""
        elapsed = time.time() - self.start_time
        name = self.getDescription(test)
        self.stream.write(f"Finished test: {name} ({elapsed:.03}s)\n")
        if name in results:
            raise AssertionError("expected all keys to be unique")
        results[name] = elapsed
        super().stopTest(test)


def print_results(results, discovery_time, thresh, status):
    # only keep results >= threshold
    results = dict(filter(lambda x: x[1] > thresh, results.items()))
    if len(results) == 0:
        return
    print(f"\n\n{status}, printing completed times >{thresh}s in ascending order...\n")
    timings = dict(sorted(results.items(), key=lambda item: item[1]))

    for r in timings:
        if timings[r] >= thresh:
            print(f"{r} ({timings[r]:.03}s)")
    print(f"test discovery time: {discovery_time:.03}s")
    print(f"total testing time: {sum(results.values()):.03}s")
    print("Remember to check above times for any errors!")


def parse_args(default_pattern):
    parser = argparse.ArgumentParser(description="Runner for MONAI unittests with timing.")
    parser.add_argument(
        "-s", action="store", dest="path", default=".", help="Directory to start discovery (default: '%(default)s')"
    )
    parser.add_argument(
        "-p",
        action="store",
        dest="pattern",
        default=default_pattern,
        help="Pattern to match tests (default: '%(default)s')",
    )
    parser.add_argument(
        "-t",
        "--thresh",
        dest="thresh",
        default=10.0,
        type=float,
        help="Display tests longer than given threshold (default: %(default)d)",
    )
    parser.add_argument(
        "-v",
        "--verbosity",
        action="store",
        dest="verbosity",
        type=int,
        default=1,
        help="Verbosity level (default: %(default)d)",
    )
    parser.add_argument("-q", "--quick", action="store_true", dest="quick", default=False, help="Only do quick tests")
    parser.add_argument(
        "-f", "--failfast", action="store_true", dest="failfast", default=False, help="Stop testing on first failure"
    )
    args = parser.parse_args()
    print(f"Running tests in folder: '{args.path}'")
    if args.pattern:
        print(f"With file pattern: '{args.pattern}'")

    return args


def get_default_pattern(loader):
    signature = inspect.signature(loader.discover)
    params = {k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty}
    return params["pattern"]


if __name__ == "__main__":

    loader = unittest.TestLoader()
    default_pattern = get_default_pattern(loader)

    # Parse input arguments
    args = parse_args(default_pattern)

    # If quick is desired, set environment variable
    if args.quick:
        os.environ["QUICKTEST"] = "True"

    # Get all test names (optionally from some path with some pattern)
    with PerfContext() as pc:
        tests = loader.discover(args.path, args.pattern)
    discovery_time = pc.total_time
    print(f"time to discover tests: {discovery_time}s")

    test_runner = unittest.runner.TextTestRunner(
        resultclass=TimeLoggingTestResult, verbosity=args.verbosity, failfast=args.failfast
    )

    # Use try catches to print the current results if encountering exception or keyboard interruption
    try:
        test_result = test_runner.run(tests)
        print_results(results, discovery_time, args.thresh, "tests finished")
        sys.exit(not test_result.wasSuccessful())
    except KeyboardInterrupt:
        print_results(results, discovery_time, args.thresh, "tests cancelled")
        sys.exit(1)
    except Exception:
        print_results(results, discovery_time, args.thresh, "exception reached")
        raise

[(66, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (107, 'inspect.signature', 'inspect.signature', False, 'import inspect\n'), (114, 'unittest.TestLoader', 'unittest.TestLoader', False, 'import unittest\n'), (130, 'unittest.runner.TextTestRunner', 'unittest.runner.TextTestRunner', False, 'import unittest\n'), (33, 'time.time', 'time.time', False, 'import time\n'), (125, 'monai.utils.PerfContext', 'PerfContext', False, 'from monai.utils import PerfContext\n'), (40, 'time.time', 'time.time', False, 'import time\n'), (141, 'sys.exit', 'sys.exit', False, 'import sys\n')]