#!/usr/bin/env python3
#
# Copyright 2016 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.
#

"""Convert a JSON descrption of a spec test into a JavaScript."""

import argparse
import io
import json
import os
import re
import struct
import sys

import find_exe
import utils
from utils import ChangeDir, ChangeExt, Error, Executable

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

F32_INF = 0x7f800000
F32_NEG_INF = 0xff800000
F32_NEG_ZERO = 0x80000000
F32_SIGN_BIT = F32_NEG_ZERO
F32_SIG_MASK = 0x7fffff
F32_QUIET_NAN = 0x7fc00000
F32_QUIET_NAN_TAG = 0x400000
F64_INF = 0x7ff0000000000000
F64_NEG_INF = 0xfff0000000000000
F64_NEG_ZERO = 0x8000000000000000
F64_SIGN_BIT = F64_NEG_ZERO
F64_SIG_MASK = 0xfffffffffffff
F64_QUIET_NAN = 0x7ff8000000000000
F64_QUIET_NAN_TAG = 0x8000000000000


def I32ToJS(value):
    # JavaScript will return all i32 values as signed.
    if value >= 2**31:
        value -= 2**32
    return str(value)


def IsNaNF32(f32_bits):
    return (F32_INF < f32_bits < F32_NEG_ZERO) or (f32_bits > F32_NEG_INF)


def ReinterpretF32(f32_bits):
    return struct.unpack('<f', struct.pack('<I', f32_bits))[0]


def NaNF32ToString(f32_bits):
    result = '-' if f32_bits & F32_SIGN_BIT else ''
    result += 'nan'
    sig = f32_bits & F32_SIG_MASK
    if sig != F32_QUIET_NAN_TAG:
        result += f':0x{sig:x}'
    return result


def F32ToWasm(f32_bits):
    if IsNaNF32(f32_bits):
        return f'f32.const {NaNF32ToString(f32_bits)}'
    elif f32_bits == F32_INF:
        return 'f32.const infinity'
    elif f32_bits == F32_NEG_INF:
        return 'f32.const -infinity'
    else:
        return f'f32.const {float.hex(ReinterpretF32(f32_bits))}'


def F32ToJS(f32_bits):
    assert not IsNaNF32(f32_bits)
    if f32_bits == F32_INF:
        return 'Infinity'
    elif f32_bits == F32_NEG_INF:
        return '-Infinity'
    else:
        return f'f32({ReinterpretF32(f32_bits)})'


def IsNaNF64(f64_bits):
    return (F64_INF < f64_bits < F64_NEG_ZERO) or (f64_bits > F64_NEG_INF)


def ReinterpretF64(f64_bits):
    return struct.unpack('<d', struct.pack('<Q', f64_bits))[0]


def NaNF64ToString(f64_bits):
    result = '-' if f64_bits & F64_SIGN_BIT else ''
    result += 'nan'
    sig = f64_bits & F64_SIG_MASK
    if sig != F64_QUIET_NAN_TAG:
        result += f':0x{sig:x}'
    return result


def F64ToWasm(f64_bits):
    if IsNaNF64(f64_bits):
        return f'f64.const {NaNF64ToString(f64_bits)}'
    elif f64_bits == F64_INF:
        return 'f64.const infinity'
    elif f64_bits == F64_NEG_INF:
        return 'f64.const -infinity'
    else:
        return f'f64.const {float.hex(ReinterpretF64(f64_bits))}'


def F64ToJS(f64_bits):
    assert not IsNaNF64(f64_bits)
    if f64_bits == F64_INF:
        return 'Infinity'
    elif f64_bits == F64_NEG_INF:
        return '-Infinity'
    else:
        # Use repr to get full precision
        return repr(ReinterpretF64(f64_bits))


def UnescapeWasmString(s):
    # Wat allows for more escape characters than this, but we assume that
    # wasm2wat will only use the \xx escapes.
    result = ''
    i = 0
    while i < len(s):
        c = s[i]
        if c == '\\':
            x = s[i + 1:i + 3]
            if len(x) != 2:
                raise Error(f'String with invalid escape: \"{s}\"')
            result += chr(int(x, 16))
            i += 3
        else:
            result += c
            i += 1
    return result


def EscapeJSString(s):
    result = ''
    for c in s:
        if 32 <= ord(c) < 127 and c not in '"\\':
            result += c
        else:
            result += f'\\x{ord(c):02x}'
    return result


def IsValidJSConstant(const):
    type_ = const['type']
    value = const['value']
    if type_ in ('f32', 'f64') and value in ('nan:canonical', 'nan:arithmetic'):
        return True
    if type_ == 'i32':
        return True
    elif type_ == 'i64':
        return False
    elif type_ == 'f32':
        return not IsNaNF32(int(value))
    elif type_ == 'f64':
        return not IsNaNF64(int(value))


def IsValidJSAction(action):
    return all(IsValidJSConstant(x) for x in action.get('args', []))


def IsValidJSCommand(command):
    type_ = command['type']
    action = command['action']
    if type_ == 'assert_return':
        expected = command['expected']
        return (IsValidJSAction(action) and
                all(IsValidJSConstant(x) for x in expected))
    elif type_ in ('assert_trap', 'assert_exhaustion'):
        return IsValidJSAction(action)


def CollectInvalidModuleCommands(commands):
    modules = []
    module_map = {}
    for command in commands:
        if command['type'] == 'module':
            pair = (command, [])
            modules.append(pair)
            module_name = command.get('name')
            if module_name:
                module_map[module_name] = pair
        elif command['type'] in ('assert_return', 'assert_trap',
                                 'assert_exhaustion'):
            if IsValidJSCommand(command):
                continue

            action = command['action']
            module_name = action.get('module')
            if module_name:
                pair = module_map[module_name]
            else:
                pair = modules[-1]
            pair[1].append(command)
    return modules


class ModuleExtender:

    def __init__(self, wat2wasm, wasm2wat, temp_dir):
        self.wat2wasm = wat2wasm
        self.wasm2wat = wasm2wat
        self.temp_dir = temp_dir
        self.lines = []
        self.exports = {}

    def Extend(self, wasm_path, commands):
        wat_path = self._RunWasm2Wat(wasm_path)
        with open(wat_path) as wat_file:
            wat = wat_file.read()

        self.lines = []
        self.exports = self._GetExports(wat)
        for i, command in enumerate(commands):
            self._Command(i, command)

        wat = wat[:wat.rindex(')')] + '\n\n'
        wat += '\n'.join(self.lines) + ')'
        # print wat
        with open(wat_path, 'w') as wat_file:
            wat_file.write(wat)
        return self._RunWat2Wasm(wat_path)

    def _Command(self, index, command):
        command_type = command['type']
        new_field = f'assert_{index}'
        if command_type == 'assert_return':
            self.lines.append(f'(func (export "{new_field}")')
            self.lines.append('block')
            self._Action(command['action'])
            for expected in command['expected']:
                self._Reinterpret(expected['type'])
                if expected['value'] in ('nan:canonical', 'nan:arithmetic'):
                    self._NanBitmask(expected['value'] == 'nan:canonical', expected['type'])
                    self._And(expected['type'])
                    self._QuietNan(expected['type'])
                else:
                    self._Constant(expected)
                    self._Reinterpret(expected['type'])
                self._Eq(expected['type'])
                self.lines.extend(['i32.eqz', 'br_if 0'])
            self.lines.extend(['return', 'end', 'unreachable', ')'])
        elif command_type in ('assert_trap', 'assert_exhaustion'):
            self.lines.append(f'(func (export "{new_field}")')
            self._Action(command['action'])
            self.lines.extend(['br 0', ')'])
        else:
            raise Error(f'Unexpected command: {command_type}')

        # Update command to point to the new exported function.
        command['action']['field'] = new_field
        command['action']['args'] = []
        command['expected'] = []

    def _GetExports(self, wat):
        result = {}
        pattern = r'^\s*\(export \"(.*?)\"\s*\((\w+) (\d+)'
        for name, type_, index in re.findall(pattern, wat, re.MULTILINE):
            result[UnescapeWasmString(name)] = (type_, index)
        return result

    def _Action(self, action):
        export = self.exports[action['field']]
        if action['type'] == 'invoke':
            for arg in action['args']:
                self._Constant(arg)
            self.lines.append(f'call {export[1]}')
        elif action['type'] == 'get':
            self.lines.append(f'global.get {export[1]}')
        else:
            raise Error(f"Unexpected action: {action['type']}")

    def _Reinterpret(self, type_):
        self.lines.extend({
            'i32': [],
            'i64': [],
            'f32': ['i32.reinterpret_f32'],
            'f64': ['i64.reinterpret_f64']
        }[type_])

    def _Eq(self, type_):
        self.lines.append({
            'i32': 'i32.eq',
            'i64': 'i64.eq',
            'f32': 'i32.eq',
            'f64': 'i64.eq'
        }[type_])

    def _And(self, type_):
        self.lines.append({
            'i32': 'i32.and',
            'i64': 'i64.and',
            'f32': 'i32.and',
            'f64': 'i64.and'
        }[type_])

    def _NanBitmask(self, canonical, type_):
        # When checking for canonical NaNs, the value can differ only in the sign
        # bit from +nan. For arithmetic NaNs, the sign bit and the rest of the tag
        # can differ as well.
        assert type_ in ('f32', 'f64')
        if not canonical:
            return self._QuietNan(type_)

        if type_ == 'f32':
            line = 'i32.const 0x7fffffff'
        else:
            line = 'i64.const 0x7fffffffffffffff'
        self.lines.append(line)

    def _QuietNan(self, type_):
        assert type_ in ('f32', 'f64')
        if type_ == 'f32':
            line = f'i32.const 0x{F32_QUIET_NAN:x}'
        else:
            line = f'i64.const 0x{F64_QUIET_NAN:x}'
        self.lines.append(line)

    def _Constant(self, const):
        inst = None
        type_ = const['type']
        value = const['value']
        assert value not in ('nan:canonical', 'nan:arithmetic')
        if type_ == 'i32':
            inst = f'i32.const {value}'
        elif type_ == 'i64':
            inst = f'i64.const {value}'
        elif type_ == 'f32':
            inst = F32ToWasm(int(value))
        elif type_ == 'f64':
            inst = F64ToWasm(int(value))
        self.lines.append(inst)

    def _RunWasm2Wat(self, wasm_path):
        wat_path = ChangeDir(ChangeExt(wasm_path, '.wat'), self.temp_dir)
        self.wasm2wat.RunWithArgs(wasm_path, '-o', wat_path)
        return wat_path

    def _RunWat2Wasm(self, wat_path):
        wasm_path = ChangeDir(ChangeExt(wat_path, '.wasm'), self.temp_dir)
        self.wat2wasm.RunWithArgs(wat_path, '-o', wasm_path)
        return wasm_path


class JSWriter:

    def __init__(self, base_dir, spec_json, out_file):
        self.base_dir = base_dir
        self.source_filename = os.path.basename(spec_json['source_filename'])
        self.commands = spec_json['commands']
        self.out_file = out_file
        self.module_idx = 0

    def Write(self):
        for command in self.commands:
            self._WriteCommand(command)

    def _WriteFileAndLine(self, command):
        line_number = command['line']
        self.out_file.write(f'// {self.source_filename}:{line_number}\n')

    def _WriteCommand(self, command):
        command_funcs = {
            'module': self._WriteModuleCommand,
            'action': self._WriteActionCommand,
            'instance': self._WriteInstanceCommand,
            'register': self._WriteRegisterCommand,
            'assert_malformed': self._WriteAssertModuleCommand,
            'assert_invalid': self._WriteAssertModuleCommand,
            'assert_unlinkable': self._WriteAssertModuleCommand,
            'assert_uninstantiable': self._WriteAssertModuleCommand,
            'assert_return': self._WriteAssertReturnCommand,
            'assert_trap': self._WriteAssertActionCommand,
            'assert_exhaustion': self._WriteAssertActionCommand,
        }

        func = command_funcs.get(command['type'])
        if func is None:
            raise Error(f"Unexpected type: {command['type']}")
        self._WriteFileAndLine(command)
        func(command)
        self.out_file.write('\n')

    def _ModuleIdxName(self):
        return f'${self.module_idx}'

    def _WriteModuleCommand(self, command):
        module = self._Module(command['filename'])
        if 'definition' in command:
            if 'name' in command:
                self.out_file.write(f"let {command['name']} = ")
            self.out_file.write(f'definition("{module}");\n')
        else:
            self.module_idx += 1
            idx_name = self._ModuleIdxName()
            self.out_file.write(f'let {idx_name} = instance("{module}");\n')
            if 'name' in command:
                self.out_file.write(f"let {command['name']} = {idx_name};\n")

    def _WriteActionCommand(self, command):
        action = self._Action(command['action'])
        self.out_file.write(f'{action};\n')

    def _WriteInstanceCommand(self, command):
        self.module_idx += 1
        idx_name = self._ModuleIdxName()
        self.out_file.write(f'let {idx_name} = instance_from({command["definition"]});\n')
        self.out_file.write(f"let {command['instance']} = {idx_name};")

    def _WriteRegisterCommand(self, command):
        name = command.get('name', self._ModuleIdxName())
        self.out_file.write(f"register(\"{command['as']}\", {name})\n")

    def _WriteAssertModuleCommand(self, command):
        # Don't bother writing out text modules; they can't be parsed by JS.
        if command['module_type'] == 'binary':
            module = self._Module(command['filename'])
            self.out_file.write(f'{command["type"]}("{module}");\n')

    def _WriteAssertReturnCommand(self, command):
        expected = command['expected']
        if len(expected) == 1:
            action = self._Action(command['action'])
            constants = self._ConstantList(expected)
            self.out_file.write(f'assert_return(() => {action}, {constants});\n')
        elif len(expected) == 0:
            self._WriteAssertActionCommand(command)
        else:
            raise Error(f'Unexpected result with multiple values: {expected}')

    def _WriteAssertActionCommand(self, command):
        action = self._Action(command['action'])
        self.out_file.write(f'{command["type"]}(() => {action});\n')

    def _Module(self, filename):
        with open(os.path.join(self.base_dir, filename), 'rb') as wasm_file:
            return ''.join(f'\\x{c:02x}' for c in bytearray(wasm_file.read()))

    def _Constant(self, const):
        assert IsValidJSConstant(const), f'Invalid JS const: {const}'
        type_ = const['type']
        value = const['value']
        if type_ in ('f32', 'f64') and value in ('nan:canonical', 'nan:arithmetic'):
            return value
        if type_ == 'i32':
            return I32ToJS(int(value))
        elif type_ == 'f32':
            return F32ToJS(int(value))
        elif type_ == 'f64':
            return F64ToJS(int(value))
        else:
            assert False

    def _ConstantList(self, consts):
        return ', '.join(self._Constant(const) for const in consts)

    def _Action(self, action):
        type_ = action['type']
        module = action.get('module', self._ModuleIdxName())
        field = EscapeJSString(action['field'])
        if type_ == 'invoke':
            args = f"[{self._ConstantList(action.get('args', []))}]"
            return f'call({module}, "{field}", {args})'
        elif type_ == 'get':
            return f'get({module}, "{field}")'
        else:
            raise Error(f'Unexpected action type: {type_}')


def main(args):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('-o', '--output', metavar='PATH', help='output file.')
    parser.add_argument('-P', '--prefix', metavar='PATH', help='prefix file.',
                        default=os.path.join(SCRIPT_DIR, 'gen-spec-prefix.js'))
    parser.add_argument('--bindir', metavar='PATH',
                        default=find_exe.GetDefaultPath(),
                        help='directory to search for all executables.')
    parser.add_argument('--temp-dir', metavar='PATH',
                        help='set the directory that temporary wasm/wat'
                        ' files are written.')
    parser.add_argument('--no-error-cmdline',
                        help='don\'t display the subprocess\'s commandline when'
                        ' an error occurs', dest='error_cmdline',
                        action='store_false')
    parser.add_argument('-p', '--print-cmd',
                        help='print the commands that are run.',
                        action='store_true')
    parser.add_argument('file', help='spec json file.')
    options = parser.parse_args(args)

    wat2wasm = Executable(
        find_exe.GetWat2WasmExecutable(options.bindir),
        error_cmdline=options.error_cmdline)
    wasm2wat = Executable(
        find_exe.GetWasm2WatExecutable(options.bindir),
        error_cmdline=options.error_cmdline)

    wat2wasm.verbose = options.print_cmd
    wasm2wat.verbose = options.print_cmd

    with open(options.file) as json_file:
        json_dir = os.path.dirname(options.file)
        spec_json = json.load(json_file)
        all_commands = spec_json['commands']

    # modules is a list of pairs: [(module_command, [assert_command, ...]), ...]
    modules = CollectInvalidModuleCommands(all_commands)

    with utils.TempDirectory(options.temp_dir, 'gen-spec-js-') as temp_dir:
        extender = ModuleExtender(wat2wasm, wasm2wat, temp_dir)
        for module_command, assert_commands in modules:
            if assert_commands:
                wasm_path = os.path.join(json_dir, module_command['filename'])
                new_module_filename = extender.Extend(wasm_path, assert_commands)
                module_command['filename'] = new_module_filename

        output = io.StringIO()
        if options.prefix:
            with open(options.prefix) as prefix_file:
                output.write(prefix_file.read())
                output.write('\n')

        JSWriter(json_dir, spec_json, output).Write()

    if options.output:
        with open(options.output, 'w') as out_file:
            out_file.write(output.getvalue())
    else:
        sys.stdout.write(output.getvalue())

    return 0


if __name__ == '__main__':
    try:
        sys.exit(main(sys.argv[1:]))
    except Error as e:
        sys.stderr.write(str(e) + '\n')
        sys.exit(1)