#!/usr/bin/env python3
"""Translate the pinned TargetDeserialiser initialization graph, without unwinding."""
import argparse
import hashlib
import json
import struct
import subprocess
from pathlib import Path
from a64_program import ProgramCompiler
from lift_rng import arm64_image, file_range

ROOT = Path(__file__).resolve().parents[1]
RANGES = [
    (0x101AB7C74, 108, '2ada8a5d403d8d30f9bc053f3cad0e74d08103d6c96bcdbf4db44d481ba6d4d0'),
    (0x10244A4BC, 1260, '0e450e3ad613dded236f6627d78a07894d48b46e845baa40877d6929043b171a'),
    (0x102FC2D98, 32, 'c991e960829a78707ab67384d3dbf5d296aaa8fc9f33878ae7c1e7571f779dce')]
HOSTS = {0x1027D36E4:'allocate', 0x1027D36B4:'unsupported_delete',
         0x10244A9D4:'unsupported_length', 0x100035550:'unsupported_array_length'}

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

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('binary',type=Path)
    args=parser.parse_args()
    image,slice_offset=arm64_image(args.binary.read_bytes())
    slice_hash=hashlib.sha256(image).hexdigest()
    if slice_hash!='2cb4de127b50d54c6c5666acfef2263fae16d11a3361ad39353ecf8a66a4d9c7':
        raise ValueError('Unrecognized Factorio ARM64 image')
    bodies={}
    for address,size,expected in RANGES:
        offset=file_range(image,address,size)
        data=image[offset:offset+size]
        if hashlib.sha256(data).hexdigest()!=expected:
            raise ValueError(f'Unrecognized target bytes at {address:#x}')
        bodies[address]=data
    compiler=ProgramCompiler(HOSTS,returning_hosts={0x1027D36E4}, memory_pages=16,
                             static_pages={0x102FC2000:65536})
    functions={address:bodies[address] for address,_,_ in RANGES[:2]}
    wat=compiler.compile(functions,RANGES[0][0]).replace('(import "env"','(import "target_host"')
    out=ROOT/'build/target-deserialiser'
    out.mkdir(parents=True,exist_ok=True)
    (out/'factorio_target.wat').write_text(wat)
    run(ROOT/'build/wabt/wat2wasm',out/'factorio_target.wat','-o',out/'factorio_target.wasm')
    run(ROOT/'build/wabt/wasm2c',out/'factorio_target.wasm','-o',out/'factorio_target.c','--module-name=factorio_target')
    # Only the desktop oracle executes copied ARM64 instructions. Relocate
    # direct calls, immediate branches and the one ADRP+ADD vtable reference.
    assembly=['.text','.p2align 2','.globl _factorio_target_original']
    labels={address:f'Ltarget_{address:x}' for address in functions}
    for address,data in functions.items():
        if address==RANGES[0][0]: assembly.append('_factorio_target_original:')
        for i,(word,) in enumerate(struct.iter_unpack('<I',data)):
            pc=address+4*i
            assembly.append(f'Ltarget_{pc:x}:')
            if word & 0xfc000000 == 0x94000000:
                target=pc+((word&0x3ffffff)-(0x4000000 if word&0x2000000 else 0))*4
                symbol=labels.get(target) or '_native_target_'+HOSTS[target]
                assembly.append('bl '+symbol)
            elif pc==0x10244A50C:
                assembly.append('adrp x9, _native_target_vtable@PAGE')
            elif pc==0x10244A510:
                assembly.append('add x9, x9, _native_target_vtable@PAGEOFF')
            else:
                assembly.append(f'.long 0x{word:08x}')
    assembly += ['.section __DATA,__const','.p2align 3','.globl _native_target_vtable',
                 '_native_target_vtable:', '.quad 0, 0, 0, 0']
    (out/'oracle.S').write_text('\n'.join(assembly)+'\n')
    runtime=ROOT/'third_party/wabt/wasm2c'
    includes=['-I'+str(p) for p in (out,ROOT/'runtime',runtime)]
    flags=['-O2','-DWASM_RT_USE_MMAP=0','-DWASM_RT_MEMCHECK_BOUNDS_CHECK=1']
    sources=[out/'factorio_target.c',ROOT/'runtime/target_bridge.c',ROOT/'runtime/guest_heap.c']
    run('xcrun','clang',*flags,*includes,*sources,runtime/'wasm-rt-impl.c',runtime/'wasm-rt-mem-impl.c',
        out/'oracle.S',ROOT/'tests/target_equivalence.c','-o',out/'verify_target')
    run(out/'verify_target')
    sdk=subprocess.check_output(['xcrun','--sdk','iphoneos','--show-sdk-path'],text=True).strip()
    objects=[]
    for source in sources:
        obj=out/(source.stem+'.ios.o')
        run('xcrun','--sdk','iphoneos','clang','-target','arm64-apple-ios15.0','-isysroot',sdk,
            *flags,*includes,'-c',source,'-o',obj)
        objects.append(obj)
    run('xcrun','libtool','-static','-o',out/'libfactorio_target_ios.a',*objects)
    manifest={'arm64_sha256':slice_hash,'ranges':[{'address':hex(a),'bytes':s,'sha256':h} for a,s,h in RANGES],
              'instructions':compiler.instructions,'calls':compiler.calls,'memory_bytes':1048576,
              'static_page':{'original':'0x102fc2000','guest':65536},
              'scope':'constructor initialization only; static vtable identity is opaque; no virtual calls, free or C++ unwinding',
              'rea_assembly_evidence':['ev_1a00dc8f1f23ce6c1e54f8d189d163fbdd012621890e0288643b8f83a8b11611',
                                       'ev_b08fa6c3b275200505c3e78c71046526444b5c0ec603ef09d18d5db53dbf0eab'],
              'rea_exact_byte_records':'validated separately by tools/check_initialization_rea.py; see rea-provenance.json',
              'excluded_unwind_ranges':['0x101ab7ce0..0x101ab7d0c','0x10244a9a8..0x10244a9d4']}
    (out/'manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
    (out/'verification.json').write_text(json.dumps({'status':'pass','native_object_graph_comparisons':512,
        'targets_per_graph':64,'allocations_per_graph':4,'heap_bytes_per_graph':2104,
        'allocation_failure_cases':8,'instructions':len(compiler.instructions),
        'scope':'normal constructor initialization and allocation aborts; no C++ unwinding'},indent=2)+'\n')
    print('PASS: original TargetDeserialiser graph, bounded heap, desktop oracle and iOS library')

if __name__=='__main__': main()