"""Execute branch/loop translations against original ARM64 instruction words."""
import json
import struct
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'tools'))
from a64_lift import Compiler, UnsupportedInstruction

OUT = ROOT / 'build/cfg'
OUT.mkdir(parents=True, exist_ok=True)
# Hand-authored instruction fixtures, NOT recovered Factorio functions.
# Original relative branch offsets are preserved in the native oracle.
FIXTURES = {
    'countdown': [0x34000060, 0x51000400, 0x17fffffe, 0xd65f03c0],
    'bit63': [0xb6f80060, 0xd2800060, 0x14000002, 0xd28000e0, 0xd65f03c0],
    'nonzero32': [0x35000060, 0xd28000a0, 0x14000002, 0xd2800160, 0xd65f03c0],
    'nonzero64': [0xb5000060, 0xd28000a0, 0x14000002, 0xd2800160, 0xd65f03c0],
    'bit0set': [0x37000060, 0xd2800060, 0x14000002, 0xd28000e0, 0xd65f03c0],
}
for operation, opcode in [('cmp32', 0x6b01001f), ('cmp64', 0xeb01001f),
                          ('cmn32', 0x2b01001f), ('cmn64', 0xab01001f)]:
    for condition in range(16):
        FIXTURES[f'{operation}c{condition}'] = [opcode, 0x54000060 | condition,
            0xd2800060, 0x14000002, 0xd28000e0, 0xd65f03c0]
for operation, opcode in [('addimm64', 0xb100041f), ('subimm32', 0x7100041f),
                          ('asr32', 0x6b817c1f), ('asr64', 0xeb81fc1f)]:
    for condition in (0, 2, 6, 10):
        FIXTURES[f'{operation}c{condition}'] = [opcode, 0x54000060 | condition,
            0xd2800060, 0x14000002, 0xd28000e0, 0xd65f03c0]
for width in (32, 64):
    for operation in range(4):
        condition = 10 if width == 32 else 8
        select = (0x1a810000 if width == 32 else 0x9a810000) | condition << 12
        select |= (operation >> 1) << 30 | (operation & 1) << 10
        FIXTURES[f'select{width}op{operation}'] = [0x14000001,
            0x6b01001f if width == 32 else 0xeb01001f, select, 0xd65f03c0]


def run(*args):
    subprocess.run(list(map(str, args)), cwd=ROOT, check=True)


assembly = ['.text', '.p2align 2']
for name, words in FIXTURES.items():
    body = struct.pack('<' + 'I' * len(words), *words)
    (OUT / f'{name}.wat').write_text(Compiler().compile(body, 0x1000))
    run(ROOT / 'build/wabt/wat2wasm', OUT / f'{name}.wat', '-o', OUT / f'{name}.wasm')
    run(ROOT / 'build/wabt/wasm2c', OUT / f'{name}.wasm', '-o', OUT / f'{name}.c', '--module-name=' + name)
    assembly += [f'.globl _native_{name}', f'_native_{name}:']
    assembly += [f'.long 0x{word:08x}' for word in words]
(OUT / 'oracle.S').write_text('\n'.join(assembly) + '\n')

for words in ([0x14000002, 0xd65f03c0],  # target outside body
              [0x14000002, 0xd2800000, 0xd65f03c0],  # unreachable code
              [0x34000040, 0x14000002, 0xd65f03c0],  # one edge leaves body
              [0x54000040, 0xd2800000, 0xd65f03c0],  # flags undefined
              [0x34000040, 0xf100041f, 0x54000040, 0xd2800000, 0xd65f03c0],  # flags defined on only one path
              [0x14000000]):  # no RET
    try:
        Compiler().compile(struct.pack('<' + 'I' * len(words), *words), 0x1000)
    except UnsupportedInstruction:
        pass
    else:
        raise AssertionError('Invalid CFG accepted')

includes = '\n'.join(f'#include "{name}.h"\nextern uint64_t native_{name}(uint64_t, uint64_t);' for name in FIXTURES)
checks = []
for name in FIXTURES:
    inputs = 'i % 201' if name == 'countdown' else 'values[i % 8] ^ (i < 64 ? 0 : random_value(&seed))'
    checks.append(f'''
  {{
    w2c_{name} module;
    wasm2c_{name}_instantiate(&module);
    uint8_t *memory = w2c_{name}_memory(&module)->data;
    for (unsigned i = 0; i < 20000; ++i) {{
      uint64_t input = {inputs};
      uint64_t registers[32];
      for (int r = 0; r < 32; ++r) registers[r] = random_value(&seed);
      registers[0] = input;
      registers[1] = i < 64 ? values[(i / 8) % 8] : registers[1];
      memset(memory, 0xa5, 65536);
      memcpy(memory, registers, sizeof(registers));
      assert(w2c_{name}_run(&module, 0, 1000) == 0);
      uint64_t output[32]; memcpy(output, memory, sizeof(output));
      assert(output[0] == native_{name}(input, registers[1]));
      assert(memcmp(output + 1, registers + 1, 31 * sizeof(uint64_t)) == 0);
      for (unsigned j = 256; j < 65536; ++j) assert(memory[j] == 0xa5);
    }}
    wasm2c_{name}_free(&module);
  }}''')
(OUT / 'verify.c').write_text(f'''
#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
{includes}
static uint64_t random_value(uint64_t *state) {{
  *state ^= *state << 13; *state ^= *state >> 7; *state ^= *state << 17; return *state;
}}
int main(void) {{
  wasm_rt_init();
  uint64_t seed = 0x123456789abcdef;
  uint64_t values[] = {{0, 1, 0xffffffff, 0x100000000, 0x8000000000000000,
      0xffffffffffffffff, 0x7fffffffffffffff, 2}};
  {''.join(checks)}
  w2c_countdown loop;
  wasm2c_countdown_instantiate(&loop);
  uint64_t start = 100;
  memcpy(w2c_countdown_memory(&loop)->data, &start, 8);
  assert(w2c_countdown_run(&loop, 0, 0) == 3);
  assert(w2c_countdown_run(&loop, 0, 1) == 3);
  assert(w2c_countdown_run(&loop, 0, 1000) == 0);
  wasm2c_countdown_free(&loop);
  wasm_rt_free();
  puts("PASS: {len(FIXTURES) * 20000} native branch/loop/NZCV comparisons, preserved registers, memory guards, budget and CFG rejection");
  return 0;
}}
''')
runtime = ROOT / 'third_party/wabt/wasm2c'
run('xcrun', 'clang', '-O2', '-DWASM_RT_USE_MMAP=0', '-DWASM_RT_MEMCHECK_BOUNDS_CHECK=1',
    '-I' + str(OUT), '-I' + str(runtime), *(OUT / f'{name}.c' for name in FIXTURES),
    OUT / 'verify.c', OUT / 'oracle.S', runtime / 'wasm-rt-impl.c', runtime / 'wasm-rt-mem-impl.c',
    '-o', OUT / 'verify')
run(OUT / 'verify')
(OUT / 'verification.json').write_text(json.dumps({'status': 'pass', 'native_comparisons': len(FIXTURES) * 20000,
    'fixtures': list(FIXTURES), 'source': 'hand-authored instruction fixtures; not game code',
    'cfg_invalid_cases': 6, 'budget_exhaustion_status': 3}, indent=2) + '\n')