# Copyright 2015 WebAssembly Community Group participants
#
# 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 difflib
import fnmatch
import glob
import io
import os
import shutil
import stat
import subprocess
import sys
from contextlib import contextmanager
from multiprocessing.pool import ThreadPool
from pathlib import Path

# The C++ standard whose features are required to build Binaryen.
# Keep in sync with CMakeLists.txt CXX_STANDARD
cxx_standard = 20


def parse_args(args):
    usage_str = ("usage: 'python check.py [options]'\n\n"
                 "Runs the Binaryen test suite.")
    parser = argparse.ArgumentParser(description=usage_str)
    parser.add_argument(
        '--torture', dest='torture', action='store_true', default=True,
        help='Chooses whether to run the torture testcases. Default: true.')
    parser.add_argument(
        '--no-torture', dest='torture', action='store_false',
        help='Disables running the torture testcases.')
    parser.add_argument(
        '--abort-on-first-failure', '--fail-fast', dest='abort_on_first_failure',
        action=argparse.BooleanOptionalAction, default=True,
        help=('Specifies whether to halt test suite execution on first test error.'
              ' Default: true.'))
    parser.add_argument(
        '--binaryen-bin', dest='binaryen_bin', default='',
        help=('Specifies the path to the Binaryen executables in the CMake build'
              ' directory. Default: bin/ of current directory (i.e. assume an'
              ' in-tree build).'
              ' If not specified, the environment variable BINARYEN_BIN= can also'
              ' be used to adjust this.'))
    parser.add_argument(
        '--binaryen-lib', dest='binaryen_lib', default='',
        help=('Specifies a path to where the built Binaryen shared library resides at.'
              ' Default: ./lib relative to bin specified above.'))
    parser.add_argument(
        '--binaryen-root', dest='binaryen_root', default='',
        help=('Specifies a path to the root of the Binaryen repository tree.'
              ' Default: the directory where this file check.py resides.'))
    parser.add_argument(
        '--out-dir', dest='out_dir', default='',
        help=('Specifies a path to the output directory for temp files, which '
              'is also where the test runner changes directory into.'
              ' Default:. out/test under the binaryen root.'))
    parser.add_argument(
        '--valgrind', dest='valgrind', default='',
        help=('Specifies a path to Valgrind tool, which will be used to validate'
              ' execution if specified. (Pass --valgrind=valgrind to search in'
              ' PATH)'))
    parser.add_argument(
        '--valgrind-full-leak-check', dest='valgrind_full_leak_check',
        action='store_true', default=False,
        help=('If specified, all unfreed (but still referenced) pointers at the'
              ' end of execution are considered memory leaks. Default: disabled.'))
    parser.add_argument(
        '--spec-test', action='append', default=[], dest='spec_tests',
        help='Names specific spec tests to run.')
    parser.add_argument(
        'positional_args', metavar='TEST_SUITE', nargs='*',
        help=('Names specific test suites to run. Use --list-suites to see a '
              'list of all test suites'))
    parser.add_argument(
        '--list-suites', action='store_true',
        help='List the test suites that can be run.')
    parser.add_argument(
        '--filter', dest='test_name_filter', default='',
        help=('Specifies a filter. Only tests whose paths contains this '
              'substring will be run'))
    # This option is only for fuzz_opt.py
    # TODO Allow each script to inherit the default set of options and add its
    # own custom options on top of that
    parser.add_argument(
        '--no-auto-initial-contents', dest='auto_initial_contents',
        action='store_false', default=True,
        help='Disables the automatic selection of important initial contents '
             'in fuzzer.')
    parser.add_argument(
        '--verbose', action='store_true', default=False,
        help='Enables verbose logging.')

    return parser.parse_args(args)


options = parse_args(sys.argv[1:])
requested = options.positional_args
script_dir = os.path.dirname(os.path.abspath(__file__))

num_failures = 0
warnings = []


def warn(text):
    warnings.append(text)
    print('warning:', text, file=sys.stderr)


def print_heading(msg):
    print(f'[ {msg} ]')


def verbose_log(*args, **kwargs):
    if options.verbose:
        print(*args, **kwargs)


@contextmanager
def red_output(file=sys.stderr):
    print("\033[31m", end="", file=file)
    try:
        yield
    finally:
        print("\033[0m", end="", file=file)


def red_stderr():
    return red_output(file=sys.stderr)


# Adapted from Emscripten's tools/utils.py
def get_num_cores():
    # Prefer `os.process_cpu_count` when available (3.13 and above) since it
    # takes into account thread affinity. Fall back to `os.sched_getaffinity`
    # where available and finally `os.cpu_count`, which should work everywhere.
    if hasattr(os, 'process_cpu_count'):
        cpu_count = os.process_cpu_count()
    elif hasattr(os, 'sched_getaffinity'):
        cpu_count = len(os.sched_getaffinity(0))
    else:
        cpu_count = os.cpu_count()
    return int(os.getenv('BINARYEN_CORES', cpu_count))


def run_parallel_tests(run_one_test_func, tests, show_worker_count=True):
    global num_failures
    tests = list(tests)
    if not tests:
        return
    worker_count = min(get_num_cores() or 1, len(tests))
    if show_worker_count:
        print(f"Running with {worker_count} workers")

    def run_test_with_wrapped_stdout(test):
        out = io.StringIO()
        try:
            run_one_test_func(test, stdout=out)
        except Exception as e:
            print(e, file=out)
            return False, out.getvalue()
        return True, out.getvalue()

    failed_stdouts = []
    with ThreadPool(processes=worker_count) as pool:
        try:
            for success, stdout in pool.imap_unordered(run_test_with_wrapped_stdout, tests):
                if success:
                    print(stdout, end="")
                    continue

                num_failures += 1
                failed_stdouts.append(stdout)
                if options.abort_on_first_failure:
                    with red_stderr():
                        print("Aborted test execution after first failure. Set --no-fail-fast to disable this.", file=sys.stderr)
                    break
        except KeyboardInterrupt:
            # Hard exit to avoid threads continuing to run after Ctrl-C.
            # There's no concern of deadlocking during shutdown here.
            os._exit(1)

    if failed_stdouts:
        with red_stderr():
            print("Failed tests:", file=sys.stderr)
            for failed in failed_stdouts:
                print(failed, end="", file=sys.stderr)


# setup

# Locate Binaryen source directory if not specified.
if not options.binaryen_root:
    options.binaryen_root = os.path.dirname(os.path.dirname(script_dir))

# Locate Binaryen build artifacts directory (bin/ by default)
if not options.binaryen_bin:
    default_bin = os.path.join(options.binaryen_root, 'bin')
    options.binaryen_bin = os.environ.get('BINARYEN_BIN', default_bin)

options.binaryen_bin = os.path.normpath(os.path.abspath(options.binaryen_bin))

if not options.binaryen_lib:
    options.binaryen_lib = os.path.join(os.path.dirname(options.binaryen_bin), 'lib')

options.binaryen_lib = os.path.normpath(os.path.abspath(options.binaryen_lib))

options.binaryen_build = os.path.dirname(options.binaryen_bin)

# ensure BINARYEN_BIN is set up
os.environ['BINARYEN_BIN'] = options.binaryen_bin

wasm_dis_filenames = ['wasm-dis', 'wasm-dis.exe', 'wasm-dis.js']
if not any(os.path.isfile(os.path.join(options.binaryen_bin, f))
           for f in wasm_dis_filenames):
    warn('Binaryen not found (or has not been successfully built to bin/ ?')

options.binaryen_test = os.path.join(options.binaryen_root, 'test')

if not options.out_dir:
    default_out_dir = os.path.join(options.binaryen_root, 'out', 'test')
    options.out_dir = os.environ.get('BINARYEN_OUT_DIR', default_out_dir)

if not os.path.exists(options.out_dir):
    os.makedirs(options.out_dir)


# Finds the given executable 'program' in PATH.
# Operates like the Unix tool 'which'.
def which(program):
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)
    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        paths = [
            # Prefer tools installed using third_party/setup.py
            os.path.join(options.binaryen_root, 'third_party', 'mozjs'),
            os.path.join(options.binaryen_root, 'third_party', 'v8'),
            os.path.join(options.binaryen_root, 'third_party', 'wabt', 'bin'),
        ] + os.environ['PATH'].split(os.pathsep)
        for path in paths:
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file
            if '.' not in fname:
                if is_exe(exe_file + '.exe'):
                    return exe_file + '.exe'
                if is_exe(exe_file + '.cmd'):
                    return exe_file + '.cmd'
                if is_exe(exe_file + '.bat'):
                    return exe_file + '.bat'


NATIVECC = (os.environ.get('CC') or which('mingw32-gcc') or
            which('gcc') or which('clang'))
NATIVEXX = (os.environ.get('CXX') or which('mingw32-g++') or
            which('g++') or which('clang++'))
NODEJS = os.environ.get('NODE') or which('node') or which('nodejs')
MOZJS = which('mozjs') or which('spidermonkey')

V8 = os.environ.get('V8') or which('v8') or which('d8')

BINARYEN_INSTALL_DIR = os.path.dirname(options.binaryen_bin)
WASM_OPT = [os.path.join(options.binaryen_bin, 'wasm-opt')]
WASM_AS = [os.path.join(options.binaryen_bin, 'wasm-as')]
WASM_DIS = [os.path.join(options.binaryen_bin, 'wasm-dis')]
WASM2JS = [os.path.join(options.binaryen_bin, 'wasm2js')]
WASM_CTOR_EVAL = [os.path.join(options.binaryen_bin, 'wasm-ctor-eval')]
WASM_SHELL = [os.path.join(options.binaryen_bin, 'wasm-shell')]
WASM_REDUCE = [os.path.join(options.binaryen_bin, 'wasm-reduce')]
WASM_EMBED = [os.path.join(options.binaryen_bin, 'wasm-embed')]
WASM_METADCE = [os.path.join(options.binaryen_bin, 'wasm-metadce')]
WASM_EMSCRIPTEN_FINALIZE = [os.path.join(options.binaryen_bin,
                                         'wasm-emscripten-finalize')]
BINARYEN_JS = os.path.join(options.binaryen_bin, 'binaryen_js.js')


def wrap_with_valgrind(cmd):
    # Exit code 97 is arbitrary, used to easily detect when an error occurs that
    # is detected by Valgrind.
    valgrind = [options.valgrind, '--quiet', '--error-exitcode=97']
    if options.valgrind_full_leak_check:
        valgrind += ['--leak-check=full', '--show-leak-kinds=all']
    return valgrind + cmd


if options.valgrind:
    WASM_OPT = wrap_with_valgrind(WASM_OPT)
    WASM_AS = wrap_with_valgrind(WASM_AS)
    WASM_DIS = wrap_with_valgrind(WASM_DIS)
    WASM_SHELL = wrap_with_valgrind(WASM_SHELL)


def in_binaryen(*args):
    return os.path.join(options.binaryen_root, *args)


os.environ['BINARYEN'] = in_binaryen()


def get_platform():
    return {'linux': 'linux',
            'linux2': 'linux',
            'darwin': 'mac',
            'win32': 'windows',
            'cygwin': 'windows'}[sys.platform]


def has_shell_timeout():
    return get_platform() != 'windows' and os.system('timeout 1s pwd') == 0


# Default options to pass to v8. These enable all features.
# See https://github.com/v8/v8/blob/master/src/wasm/wasm-feature-flags.h
V8_OPTS = [
    '--wasm-staging',
    '--wasm-compilation-hints',
    '--wasm-stringref',
    '--wasm-fp16',
    '--wasm-custom-descriptors',
    '--wasm-js-interop',
    '--wasm-acquire-release',
    '--wasm-wide-arithmetic',
    '--wasm-compact-imports',
]

# external tools

try:
    if NODEJS is not None:
        subprocess.run([NODEJS, '--version'], check=True, capture_output=True)
except (OSError, subprocess.CalledProcessError):
    NODEJS = None
if NODEJS is None:
    warn('no node found (did not check proper js form)')


# utilities

# removes a file if it exists, using any and all ways of doing so
def delete_from_orbit(filename):
    try:
        os.unlink(filename)
    except OSError:
        pass
    if not os.path.exists(filename):
        return
    try:
        shutil.rmtree(filename, ignore_errors=True)
    except OSError:
        pass
    if not os.path.exists(filename):
        return
    try:
        os.chmod(filename, os.stat(filename).st_mode | stat.S_IWRITE)

        def remove_readonly_and_try_again(func, path, exc_info):
            if not (os.stat(path).st_mode & stat.S_IWRITE):
                os.chmod(path, os.stat(path).st_mode | stat.S_IWRITE)
                func(path)
            else:
                raise exc_info[1]
        shutil.rmtree(filename, onerror=remove_readonly_and_try_again)
    except OSError:
        pass


def run_process(cmd, check=True, input=None, decode_output=True, *args, **kwargs):
    if input and type(input) is str:
        input = bytes(input, 'utf-8')
    ret = subprocess.run(cmd, *args, check=check, input=input, **kwargs)
    if decode_output and ret.stdout is not None:
        ret.stdout = ret.stdout.decode('utf-8')
    if ret.stderr is not None:
        ret.stderr = ret.stderr.decode('utf-8')
    return ret


def fail_with_error(msg):
    global num_failures
    try:
        num_failures += 1
        raise Exception(msg)
    except Exception as e:
        print(str(e))
        if options.abort_on_first_failure:
            raise


def fail(actual, expected, fromfile='expected'):
    diff_lines = difflib.unified_diff(
        expected.split('\n'), actual.split('\n'),
        fromfile=fromfile, tofile='actual')
    diff_str = ''.join([a.rstrip() + '\n' for a in diff_lines])[:]
    fail_with_error(f'incorrect output, diff:\n\n{diff_str}')


def fail_if_not_identical(actual, expected, fromfile='expected'):
    if expected != actual:
        fail(actual, expected, fromfile=fromfile)


def fail_if_not_contained(actual, expected):
    if expected not in actual:
        fail(actual, expected)


def fail_if_not_identical_to_file(actual, expected_file):
    binary = expected_file.endswith(".wasm") or type(actual) is bytes
    with open(expected_file, 'rb' if binary else 'r') as f:
        fail_if_not_identical(actual, f.read(), fromfile=expected_file)


def get_test_dir(name):
    """Return the test directory located at <binaryen_root>/test/[name]."""
    return os.path.join(options.binaryen_test, name)


def get_tests(test_dir, extensions=[], recursive=False):
    """Return the list of test files in a given directory.

    'extensions' is a list of file extensions. If 'extensions' is empty, returns all files.
    """
    tests = []
    star = '**/*' if recursive else '*'
    if not extensions:
        tests += glob.glob(os.path.join(test_dir, star), recursive=True)
    for ext in extensions:
        tests += glob.glob(os.path.join(test_dir, star + ext), recursive=True)
    if options.test_name_filter:
        tests = fnmatch.filter(tests, options.test_name_filter)
    tests = [item for item in tests if os.path.isfile(item)]
    return sorted(tests)


if options.spec_tests:
    non_existent_tests = [test_name for test_name in options.spec_tests if not os.path.isfile(test_name)]
    if non_existent_tests:
        raise ValueError(f"Supplied test files do not exist: {non_existent_tests}")
    options.spec_tests = [os.path.abspath(t) for t in options.spec_tests]
else:
    options.spec_tests = get_tests(get_test_dir('spec'), ['.wast'], recursive=True)

os.chdir(options.out_dir)

# 11/27/2019: We updated the spec test suite to upstream spec repo. For some
# files that started failing after this update, we added the new files to this
# skip-list and preserved old ones by renaming them to 'old_[FILENAME].wast'
# not to lose coverage. When the cause of the error is fixed or the unsupported
# construct gets support so the new test passes, we can delete the
# corresponding 'old_[FILENAME].wast' file. When you fix the new file and
# delete the old file, make sure you rename the corresponding .wast.log file in
# expected-output/ if any.
# Paths are relative to the test/spec directory
SPEC_TESTS_TO_SKIP = [
    # Requires us to write our own floating point parser
    'const.wast',

    # Invalid module accepted
    'unreached-invalid.wast',

    # Test invalid
    'elem.wast',

    # Requires scoping of `register` statements within `thread` blocks
    'threads/thread.wast',

    # Requires better support for multi-threaded tests
    'threads/wait_notify.wast',
]
SPEC_TESTSUITE_PROPOSALS_TO_SKIP = [
]

# Paths are relative to the test/spec/testsuite directory
SPEC_TESTSUITE_TESTS_TO_SKIP = [
    'array_new_elem.wast',  # Failure to parse element segment item abbreviation
    'binary.wast',   # Missing data count section validation
    'comments.wast',  # Issue with carriage returns being treated as newlines
    'const.wast',    # Hex float constant not recognized as out of range
    'data.wast',    # Fail to parse data segment offset abbreviation
    'elem.wast',    # Requires modeling empty declarative segments
    'func.wast',    # Duplicate parameter names not properly rejected
    'if.wast',      # Requires more precise unreachable validation
    'imports.wast',  # Requires fixing handling of mutation to imported globals
    'proposals/threads/imports.wast',  # Missing memory type validation on instantiation
    'proposals/threads/memory.wast',  # Missing memory type validation on instantiation
    'annotations.wast',  # String annotations IDs should be allowed
    'table64.wast',   # Requires validations for table size
    'tag.wast',      # Non-empty tag results allowed by stack switching
    'local_init.wast',  # Requires local validation to respect unnamed blocks
    'ref_func.wast',   # Requires rejecting undeclared functions references
    'return_call_indirect.wast',  # Requires more precise unreachable validation
    'select.wast',  # Missing validation of type annotation on select
    'unreached-invalid.wast',  # Requires more precise unreachable validation
    'array.wast',  # Failure to parse element segment item abbreviation
    'br_if.wast',  # Requires more precise branch validation
    'br_on_cast.wast',  # Requires host references to not be externalized i31refs
    'br_on_cast_fail.wast',  # Requires host references to not be externalized i31refs
    'extern.wast',    # Requires ref.host wast constants
    'ref_cast.wast',  # Requires host references to not be externalized i31refs
    'ref_test.wast',  # Requires host references to not be externalized i31refs
    'struct.wast',    # Fails to roundtrip unnamed types e.g. `(ref 0)`
    'memory64.wast',        # Requires validations on the max memory size
    'imports3.wast',  # Requires better checking of exports from the special "spectest" module
    'relaxed_dot_product.wast',   # i16x8.relaxed_dot_i8x16_i7x16_s instruction not supported
    'relaxed_laneselect.wast',    # i8x16.relaxed_laneselect instruction not supported
    'simd_const.wast',            # Hex float constant not recognized as out of range
    'token.wast',                 # Lexer should require spaces between strings and non-paren tokens
]


def _can_run_spec_test(test):
    test = Path(test)
    if 'testsuite' not in test.parts:
        return not any(test.match(f"test/spec/{test_to_skip}") for test_to_skip in SPEC_TESTS_TO_SKIP)

    if any(proposal in test.parts for proposal in SPEC_TESTSUITE_PROPOSALS_TO_SKIP):
        return False

    return not any(Path(test).match(f"test/spec/testsuite/{test_to_skip}") for test_to_skip in SPEC_TESTSUITE_TESTS_TO_SKIP)


options.spec_tests = [t for t in options.spec_tests if _can_run_spec_test(t)]


# check utilities


def binary_format_check(wast, verify_final_result=True, base_name=None, stdout=None):
    # checks we can convert the wast to binary and back

    as_file = f"{base_name}-a.wasm" if base_name is not None else "a.wasm"
    disassembled_file = f"{base_name}-ab.wast" if base_name is not None else "ab.wast"

    verbose_log('         (binary format check)', file=stdout)
    cmd = WASM_AS + [wast, '-o', as_file, '-all', '-g']
    verbose_log('            ', ' '.join(cmd), file=stdout)
    if os.path.exists(as_file):
        os.unlink(as_file)
    subprocess.check_call(cmd, stdout=subprocess.PIPE)
    assert os.path.exists(as_file)

    cmd = WASM_DIS + [as_file, '-o', disassembled_file, '-all']
    verbose_log('            ', ' '.join(cmd), file=stdout)
    if os.path.exists(disassembled_file):
        os.unlink(disassembled_file)
    subprocess.check_call(cmd, stdout=subprocess.PIPE)
    assert os.path.exists(disassembled_file)

    # make sure it is a valid wast
    cmd = WASM_OPT + [disassembled_file, '-all', '-q']
    verbose_log('            ', ' '.join(cmd), file=stdout)
    subprocess.check_call(cmd, stdout=subprocess.PIPE)

    if verify_final_result:
        actual = open(disassembled_file).read()
        fail_if_not_identical_to_file(actual, wast + '.fromBinary')

    return disassembled_file


def pass_debug_env():
    env = os.environ.copy()
    env['BINARYEN_PASS_DEBUG'] = '1'
    return env


@contextmanager
def with_pass_debug():
    old_pass_debug = os.environ.get('BINARYEN_PASS_DEBUG')
    os.environ['BINARYEN_PASS_DEBUG'] = '1'
    try:
        yield
    finally:
        if old_pass_debug is not None:
            os.environ['BINARYEN_PASS_DEBUG'] = old_pass_debug
        elif 'BINARYEN_PASS_DEBUG' in os.environ:
            del os.environ['BINARYEN_PASS_DEBUG']


# checks if we are on windows, and if so logs out that a test is being skipped,
# and returns True. This is a central location for all test skipping on
# windows, so that we can easily find which tests are skipped.
def skip_if_on_windows(name):
    if get_platform() == 'windows':
        print(f'skipping test "{name}" on windows')
        return True
    return False


test_suffixes = ['*.wasm', '*.wast', '*.wat']


# return a list of all the tests in the entire test suite
def get_all_tests():
    core_tests = get_tests(get_test_dir('.'), test_suffixes)
    passes_tests = get_tests(get_test_dir('passes'), test_suffixes)
    spec_tests = get_tests(get_test_dir('spec'), test_suffixes)
    wasm2js_tests = get_tests(get_test_dir('wasm2js'), test_suffixes)
    lld_tests = get_tests(get_test_dir('lld'), test_suffixes)
    unit_tests = get_tests(get_test_dir(os.path.join('unit', 'input')), test_suffixes)
    lit_tests = get_tests(get_test_dir('lit'), test_suffixes, recursive=True)

    return core_tests + passes_tests + spec_tests + wasm2js_tests + lld_tests + unit_tests + lit_tests