#! /usr/bin/env python3

import os
import subprocess
import sys
import argparse
from glob import glob

# Tests require at least CPython 3.3. If your default python3 executable
# is of lower version, you can point MICROPY_CPYTHON3 environment var
# to the correct executable.
if os.name == 'nt':
    CPYTHON3 = os.getenv('MICROPY_CPYTHON3', 'python3.exe')
    MICROPYTHON = os.getenv('MICROPY_MICROPYTHON', '../windows/micropython.exe')
else:
    CPYTHON3 = os.getenv('MICROPY_CPYTHON3', 'python3')
    MICROPYTHON = os.getenv('MICROPY_MICROPYTHON', '../unix/micropython')

def rm_f(fname):
    if os.path.exists(fname):
        os.remove(fname)

def run_tests(pyb, tests, args):
    test_count = 0
    testcase_count = 0
    passed_count = 0
    failed_tests = []
    skipped_tests = []

    skip_tests = set()

    # Some tests shouldn't be run under Travis CI
    if os.getenv('TRAVIS') == 'true':
        skip_tests.add('basics/memoryerror.py')

    # Some tests are known to fail with native emitter
    # Remove them from the below when they work
    if args.emit == 'native':
        skip_tests.update({'basics/%s.py' % t for t in 'bytes class_store_class class_super class_super_object closure1 closure2 closure_defargs compare_multi containment del_deref del_local del_name fun3 fun_calldblstar fun_callstar fun_callstardblstar fun_defargs fun_defargs2 fun_kwargs fun_kwonly fun_kwonlydef fun_kwvarargs fun_varargs gen_yield_from gen_yield_from_close gen_yield_from_ducktype gen_yield_from_exc gen_yield_from_iter gen_yield_from_send gen_yield_from_throw generator1 generator2 generator_args generator_close generator_closure generator_exc generator_return generator_send globals_del is_isnot string_format string_join subclass_native2_list subclass_native2_tuple subclass_native3 try2 try_as_var try_finally1 try_finally_loops try_finally_return try_reraise try_reraise2 unary_op unboundlocal with1 with_break with_continue with_return'.split()})
        skip_tests.add('float/string_format.py')
        skip_tests.add('import/gen_context.py')
        skip_tests.add('import/import_pkg1.py')
        skip_tests.add('io/file_with.py')
        skip_tests.add('micropython/heapalloc.py')
        skip_tests.add('micropython/viper.py')
        skip_tests.add('misc/features.py')
        skip_tests.add('misc/recursion.py')
        skip_tests.add('misc/rge_sm.py')

    for test_file in tests:
        test_basename = os.path.basename(test_file)
        test_name = os.path.splitext(test_basename)[0]

        if test_file in skip_tests:
            print("skip ", test_file)
            skipped_tests.append(test_name)
            continue

        # get expected output
        test_file_expected = test_file + '.exp'
        if os.path.isfile(test_file_expected):
            # expected output given by a file, so read that in
            with open(test_file_expected, 'rb') as f:
                output_expected = f.read()
                if os.name == 'nt':
                    output_expected = output_expected.replace(b'\n', b'\r\n')
        else:
            # run CPython to work out expected output
            try:
                output_expected = subprocess.check_output([CPYTHON3, '-B', test_file])
                if args.write_exp:
                    with open(test_file_expected, 'wb') as f:
                        f.write(output_expected)
            except subprocess.CalledProcessError:
                output_expected = b'CPYTHON3 CRASH'

        if args.write_exp:
            continue

        # run Micro Python
        if pyb is None:
            # run on PC
            try:
                output_mupy = subprocess.check_output([MICROPYTHON, '-X', 'emit=' + args.emit, test_file])
            except subprocess.CalledProcessError:
                output_mupy = b'CRASH'
        else:
            # run on pyboard
            pyb.enter_raw_repl()
            try:
                output_mupy = pyb.execfile(test_file).replace(b'\r\n', b'\n')
            except pyboard.PyboardError:
                output_mupy = b'CRASH'

        if output_mupy == b'SKIP\n':
            print("skip ", test_file)
            skipped_tests.append(test_name)
            continue

        testcase_count += len(output_expected.splitlines())

        filename_expected = test_basename + ".exp"
        filename_mupy = test_basename + ".out"

        if output_expected == output_mupy:
            print("pass ", test_file)
            passed_count += 1
            rm_f(filename_expected)
            rm_f(filename_mupy)
        else:
            with open(filename_expected, "wb") as f:
                f.write(output_expected)
            with open(filename_mupy, "wb") as f:
                f.write(output_mupy)
            print("FAIL ", test_file)
            failed_tests.append(test_name)

        test_count += 1

    print("{} tests performed ({} individual testcases)".format(test_count, testcase_count))
    print("{} tests passed".format(passed_count))

    if len(skipped_tests) > 0:
        print("{} tests skipped: {}".format(len(skipped_tests), ' '.join(skipped_tests)))
    if len(failed_tests) > 0:
        print("{} tests failed: {}".format(len(failed_tests), ' '.join(failed_tests)))
        return False

    # all tests succeeded
    return True

def main():
    cmd_parser = argparse.ArgumentParser(description='Run tests for Micro Python.')
    cmd_parser.add_argument('--pyboard', action='store_true', help='run the tests on the pyboard')
    cmd_parser.add_argument('-d', '--test-dirs', nargs='*', help='input test directories (if no files given)')
    cmd_parser.add_argument('--write-exp', action='store_true', help='save .exp files to run tests w/o CPython')
    cmd_parser.add_argument('--emit', default='bytecode', help='Micro Python emitter to use (bytecode or native)')
    cmd_parser.add_argument('files', nargs='*', help='input test files')
    args = cmd_parser.parse_args()

    if args.pyboard:
        import pyboard
        pyb = pyboard.Pyboard('/dev/ttyACM0')
        pyb.enter_raw_repl()
    else:
        pyb = None

    if len(args.files) == 0:
        if args.test_dirs is None:
            if pyb is None:
                # run PC tests
                test_dirs = ('basics', 'micropython', 'float', 'import', 'io', 'misc', 'unicode')
            else:
                # run pyboard tests
                test_dirs = ('basics', 'micropython', 'float', 'pyb', 'pybnative', 'inlineasm')
        else:
            # run tests from these directories
            test_dirs = args.test_dirs
        tests = sorted(test_file for test_files in (glob('{}/*.py'.format(dir)) for dir in test_dirs) for test_file in test_files)
    else:
        # tests explicitly given
        tests = args.files

    if not run_tests(pyb, tests, args):
        sys.exit(1)

if __name__ == "__main__":
    main()