#!/usr/bin/env python3

import argparse
import os
import re
import subprocess
import sys
import time
import math

parser = argparse.ArgumentParser(description='Run tests.')
parser.add_argument('--verbose', action='store_true',
                    help = 'Print failure information')
parser.add_argument('--benchmark', action='store_true',
                    help = 'Output time take to run tests')
parser.add_argument('--dir', metavar='TEST_DIR', action='append', dest='dirs',
                    help='directory containing tests to run')
parser.add_argument('--command', metavar='CMD', action='append', dest='commands',
                    help='python interpreter command')
parser.add_argument('--repeat', metavar='TIMES', type=int, default=1,
                    help='run benchmarks multiple times')
args = parser.parse_args()

if not args.commands:
    sys.exit("No commands to test")

if not args.dirs:
    baseDir = os.path.dirname(os.path.dirname(sys.argv[0]))
    args.dirs = []
    if not args.benchmark:
        args.dirs.append(os.path.join(baseDir, 'test', 'tests'))
    args.dirs.append(os.path.join(baseDir, 'bench', 'tests'))

runCount = 0
passCount = 0

def runTest(filename, command):
    global runCount, passCount

    testArgs = []
    testExpect = []
    benchArgs = []
    benchExpect = []
    benchOnly = False

    with open(filename) as f:
        while True:
            line = f.readline().rstrip()
            match = re.match('# ([\w-]+)(?:: (.*)\s*)?$', line)
            if not match:
                break

            key, value = match.groups()
            if key == 'output':
                testExpect.append(value)
            elif key == 'bench-output':
                benchExpect.append(value)
            elif key == 'args':
                testArgs = value.split()
            elif key == 'bench-args':
                benchArgs = value.split()
            elif key == 'bench-only':
                benchOnly = True
            else:
                return 'BAD', 'Error parsing header option: ' + key

    if args.benchmark:
        if benchArgs:
            testArgs = benchArgs
        if benchExpect:
            testExpect = benchExpect
    else:
        if benchOnly:
            return 'SKIP', 'Test skipped'

    runCount += 1

    if not testExpect:
        return 'BAD', 'No expected output header'

    command = command.split() + [filename] + testArgs

    runTimes = []
    for i in range(0, args.repeat):
        ok = True
        startTime = time.perf_counter()
        try:
            output = subprocess.check_output(command, stderr = subprocess.STDOUT)
            endTime = time.perf_counter()
        except subprocess.CalledProcessError as e:
            ok = False
            output = e.output

        output = output.decode().strip()

        if not ok:
            return 'FAIL', 'Execution failed: ' + output

        if not args.benchmark:
            output = output.strip()
            if output != "\n".join(testExpect):
                message = ('    Output did not match expected:\n' +
                           '      Expected: ' + '|'.join(testExpect) + '\n' +
                           '      Actual:   ' + output.replace('\n', '|'))
                return 'FAIL', message

        runTimes.append(endTime - startTime)

    passCount += 1
    return 'PASS', runTimes

def calculateTimeStats(runTimes):
    n = len(runTimes)
    mean = sum(runTimes) / n
    if n == 1:
        return mean, 0

    stdDev = math.sqrt(sum(map(lambda x: (x - mean) ** 2, runTimes)) / (n - 1))
    return mean, stdDev

def formatTimes(meanSeconds, stdDev):
    if stdDev == 0:
        return "%-13.2f       " % (meanSeconds * 1000)
    return "%-13.2f %-6.2f" % (meanSeconds * 1000, stdDev / meanSeconds * 100)

def formatResult(result):
    return "%-20s" % result

def formatScore(score):
    return "%-20.2f" % score

def formatExeName(command):
    exeName = os.path.basename(command.split()[0])
    return formatResult(exeName)

def runTestDir(dirname, names):
    commandCount = len(args.commands)
    scores = [1] * commandCount
    for name in names:
        if name.endswith('.py'):
            filename = os.path.join(dirname, name)
            results = list(map(lambda c: runTest(filename, c), args.commands))
            line = '  %-20s ' % name
            for index in range(0, commandCount):
                result, extra  = results[index]
                if result == 'SKIP':
                    line += formatResult("skipped")
                elif result != 'PASS':
                    line += formatResult(result)
                elif args.benchmark:
                    meanSeconds, stdDev = calculateTimeStats(extra)
                    scores[index] *= meanSeconds
                    line += formatTimes(meanSeconds, stdDev)
                else:
                    line += formatResult("ok")
            print(line)

            if args.verbose:
                for result, extra in results:
                    if result != 'PASS' and result != 'SKIP':
                        print(extra)

    if args.benchmark and passCount == runCount:
        line = '%-22s ' % 'Score:'
        for score in scores:
            score = 1000 * score ** (1 / runCount)
            line += formatScore(score)
        print(line)

line = 'Running tests:         '
line += "".join(map(formatExeName, args.commands))
print(line)

try:
    for testDir in args.dirs:
        for dirpath, dirnames, filenames in os.walk(testDir):
            runTestDir(dirpath, filenames)
    print("%d tests passed out of %d" % (passCount, runCount))
    if passCount != runCount:
        sys.exit(1)
except KeyboardInterrupt:
    print("Interrupted!")
    sys.exit(2)
