#! /usr/bin/env python3

import os
import subprocess
import sys
import platform
import argparse
import re
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')

# Set PYTHONIOENCODING so that CPython will use utf-8 on systems which set another encoding in the locale
os.environ['PYTHONIOENCODING'] = 'utf-8'
os.environ['MICROPYPATH'] = ''

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

def run_micropython(pyb, args, test_file):
    if pyb is None:
        # run on PC
        if test_file.startswith(('cmdline/', 'feature_check/')) or test_file == 'micropython/meminfo.py':
            # special handling for tests of the unix cmdline program

            # check for any cmdline options needed for this test
            args = [MICROPYTHON]
            with open(test_file, 'rb') as f:
                line = f.readline()
                if line.startswith(b'# cmdline:'):
                    # subprocess.check_output on Windows only accepts strings, not bytes
                    args += [str(c, 'utf-8') for c in line[10:].strip().split()]

            # run the test, possibly with redirected input
            try:
                if 'repl_' in test_file:
                    # Need to use a PTY to test command line editing
                    try:
                        import pty
                    except ImportError:
                        # in case pty module is not available, like on Windows
                        return b'SKIP\n'
                    import select

                    def get():
                        rv = b''
                        while True:
                            ready = select.select([master], [], [], 0.02)
                            if ready[0] == [master]:
                                rv += os.read(master, 1024)
                            else:
                                return rv

                    def send_get(what):
                        os.write(master, what)
                        return get()

                    with open(test_file, 'rb') as f:
                        # instead of: output_mupy = subprocess.check_output(args, stdin=f)
                        master, slave = pty.openpty()
                        p = subprocess.Popen(args, stdin=slave, stdout=slave,
                                             stderr=subprocess.STDOUT, bufsize=0)
                        banner = get()
                        output_mupy = banner + b''.join(send_get(line) for line in f)
                        p.kill()
                        os.close(master)
                        os.close(slave)
                else:
                    output_mupy = subprocess.check_output(args + [test_file])
            except subprocess.CalledProcessError:
                return b'CRASH'

            # unescape wanted regex chars and escape unwanted ones
            def convert_regex_escapes(line):
                cs = []
                escape = False
                for c in str(line, 'utf8'):
                    if escape:
                        escape = False
                        cs.append(c)
                    elif c == '\\':
                        escape = True
                    elif c in ('(', ')', '[', ']', '{', '}', '.', '*', '+', '^', '$'):
                        cs.append('\\' + c)
                    else:
                        cs.append(c)
                # accept carriage-return(s) before final newline
                if cs[-1] == '\n':
                    cs[-1] = '\r*\n'
                return bytes(''.join(cs), 'utf8')

            # convert parts of the output that are not stable across runs
            with open(test_file + '.exp', 'rb') as f:
                lines_exp = []
                for line in f.readlines():
                    if line == b'########\n':
                        line = (line,)
                    else:
                        line = (line, re.compile(convert_regex_escapes(line)))
                    lines_exp.append(line)
            lines_mupy = [line + b'\n' for line in output_mupy.split(b'\n')]
            if output_mupy.endswith(b'\n'):
                lines_mupy = lines_mupy[:-1] # remove erroneous last empty line
            i_mupy = 0
            for i in range(len(lines_exp)):
                if lines_exp[i][0] == b'########\n':
                    # 8x #'s means match 0 or more whole lines
                    line_exp = lines_exp[i + 1]
                    skip = 0
                    while i_mupy + skip < len(lines_mupy) and not line_exp[1].match(lines_mupy[i_mupy + skip]):
                        skip += 1
                    if i_mupy + skip >= len(lines_mupy):
                        lines_mupy[i_mupy] = b'######## FAIL\n'
                        break
                    del lines_mupy[i_mupy:i_mupy + skip]
                    lines_mupy.insert(i_mupy, b'########\n')
                    i_mupy += 1
                else:
                    # a regex
                    if lines_exp[i][1].match(lines_mupy[i_mupy]):
                        lines_mupy[i_mupy] = lines_exp[i][0]
                    else:
                        #print("don't match: %r %s" % (lines_exp[i][1], lines_mupy[i_mupy])) # DEBUG
                        pass
                    i_mupy += 1
                if i_mupy >= len(lines_mupy):
                    break
            output_mupy = b''.join(lines_mupy)

        else:
            # a standard test
            try:
                output_mupy = subprocess.check_output([MICROPYTHON, '-X', 'emit=' + args.emit, test_file])
            except subprocess.CalledProcessError:
                output_mupy = b'CRASH'
    else:
        # run on pyboard
        import 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'

    return output_mupy

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

    skip_tests = set()
    skip_native = False

    # Check if micropython.native is supported, and skip such tests if it's not
    native = run_micropython(pyb, args, 'feature_check/native_check.py')
    if native == b'CRASH':
        skip_native = True

    # Check if emacs repl is supported, and skip such tests if it's not
    t = run_micropython(pyb, args, 'feature_check/repl_emacs_check.py')
    if not 'True' in str(t, 'ascii'):
        skip_tests.add('cmdline/repl_emacs_keys.py')

    upy_byteorder = run_micropython(pyb, args, 'feature_check/byteorder.py')
    cpy_byteorder = subprocess.check_output([CPYTHON3, 'feature_check/byteorder.py'])
    skip_endian = (upy_byteorder != cpy_byteorder)

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

    # Some tests shouldn't be run on pyboard
    if pyb is not None:
        skip_tests.add('float/float_divmod.py') # tested by float/float_divmod_relaxed.py instead
        skip_tests.add('float/float2int_doubleprec.py') # requires double precision floating point to work
        skip_tests.add('micropython/meminfo.py') # output is very different to PC output

        if args.target == 'wipy':
            skip_tests.add('misc/print_exception.py')       # requires error reporting full
            skip_tests.add('misc/recursion.py')             # requires stack checking enabled
            skip_tests.add('misc/recursive_data.py')        # requires stack checking enabled
            skip_tests.add('misc/recursive_iternext.py')    # requires stack checking enabled
            skip_tests.add('misc/rge_sm.py')                # requires floating point
            skip_tests.update({'extmod/uctypes_%s.py' % t for t in 'bytearray le native_le ptr_le ptr_native_le sizeof sizeof_native'.split()}) # requires uctypes
            skip_tests.add('extmod/zlibd_decompress.py')    # requires zlib
            skip_tests.add('extmod/ujson_dumps_float.py')   # requires floating point
            skip_tests.add('extmod/ujson_loads_float.py')   # requires floating point
            skip_tests.add('extmod/machine1.py')            # raw memory access not supported by WiPy
            skip_tests.add('extmod/machine_mem.py')         # raw memory access not supported by WiPy
            skip_tests.add('extmod/uheapq1.py')             # uheapq not supported by WiPy
            skip_tests.add('basics/exception_chain.py')     # warning is not printed

    # Some tests are known to fail on 64-bit machines
    if pyb is None and platform.architecture()[0] == '64bit':
        pass

    # Some tests use unsupported features on Windows
    if os.name == 'nt':
        skip_tests.add('import/import_file.py') # works but CPython prints forward slashes

    # 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 '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'.split()}) # require yield
        skip_tests.update({'basics/%s.py' % t for t in 'bytes_gen class_store_class globals_del string_join'.split()}) # require yield
        skip_tests.update({'basics/%s.py' % t for t in 'try_reraise try_reraise2'.split()}) # require raise_varargs
        skip_tests.update({'basics/%s.py' % t for t in 'with1 with_break with_continue with_return'.split()}) # require with
        skip_tests.add('basics/array_construct2.py') # requires generators
        skip_tests.add('basics/bool1.py') # seems to randomly fail
        skip_tests.add('basics/del_deref.py') # requires checking for unbound local
        skip_tests.add('basics/del_local.py') # requires checking for unbound local
        skip_tests.add('basics/exception_chain.py') # raise from is not supported
        skip_tests.add('basics/try_finally_loops.py') # requires proper try finally code
        skip_tests.add('basics/try_finally_return.py') # requires proper try finally code
        skip_tests.add('basics/unboundlocal.py') # requires checking for unbound local
        skip_tests.add('import/gen_context.py') # requires yield_value
        skip_tests.add('io/file_with.py') # requires with
        skip_tests.add('io/stringio_with.py') # requires with
        skip_tests.add('misc/features.py') # requires raise_varargs
        skip_tests.add('misc/rge_sm.py') # requires yield
        skip_tests.add('misc/print_exception.py') # because native doesn't have proper traceback info
        skip_tests.add('misc/sys_exc_info.py') # sys.exc_info() is not supported for native

    for test_file in tests:
        test_file = test_file.replace('\\', '/')
        test_basename = os.path.basename(test_file)
        test_name = os.path.splitext(test_basename)[0]
        is_native = test_name.startswith("native_") or test_name.startswith("viper_")
        is_endian = test_name.endswith("_endian")

        if test_file in skip_tests or (skip_native and is_native) or (skip_endian and is_endian):
            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
        output_mupy = run_micropython(pyb, args, test_file)
        if os.name != 'nt':
            # It may be the case that we run Windows build under Linux
            # (using Wine).
            output_mupy = output_mupy.replace(b'\r\n', b'\n')

        if output_mupy == b'SKIP\n' or output_mupy == b'SKIP\r\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 MicroPython.')
    cmd_parser.add_argument('--target', default='unix', help='the target platform')
    cmd_parser.add_argument('--device', default='/dev/ttyACM0', help='the serial device or the IP address of the pyboard')
    cmd_parser.add_argument('-b', '--baudrate', default=115200, help='the baud rate of the serial device')
    cmd_parser.add_argument('-u', '--user', default='micro', help='the telnet login username')
    cmd_parser.add_argument('-p', '--password', default='python', help='the telnet login password')
    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='MicroPython emitter to use (bytecode or native)')
    cmd_parser.add_argument('files', nargs='*', help='input test files')
    args = cmd_parser.parse_args()

    if args.target == 'pyboard' or args.target == 'wipy':
        import pyboard
        pyb = pyboard.Pyboard(args.device, args.baudrate, args.user, args.password)
        pyb.enter_raw_repl()
    elif args.target == 'unix':
        pyb = None
    else:
        raise ValueError('target must be either unix, pyboard or wipy')

    if len(args.files) == 0:
        if args.test_dirs is None:
            if args.target == 'pyboard':
                # run pyboard tests
                test_dirs = ('basics', 'micropython', 'float', 'misc', 'extmod', 'pyb', 'pybnative', 'inlineasm')
            elif args.target == 'wipy':
                # run WiPy tests
                test_dirs = ('basics', 'micropython', 'misc', 'extmod', 'wipy')
            else:
                # run PC tests
                test_dirs = ('basics', 'micropython', 'float', 'import', 'io', 'misc', 'unicode', 'extmod', 'unix', 'cmdline')
        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()