#!/usr/bin/env python3
"""Lift original MapDeserialiser initialization and its actual C1 target dependency."""
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
from build_target_deserialiser import RANGES as TARGET_RANGES, HOSTS as TARGET_HOSTS
ROOT=Path(__file__).resolve().parents[1]
ENTRY=0x1014DEF40
RANGES=[(ENTRY,1272,'afef8906673fbfdab8f47b5663112ce9c4666a757bd1c80ab7e670e121f383ea'),
        (0x1012C1044,116,'c607983877925865ddb27bca5841f45a88f90896250f36864deac97462687d08'),
        *TARGET_RANGES[1:]]
HOSTS={**TARGET_HOSTS,0x10035E264:'read_or_throw',0x1027D3918:'zero_memory',
       0x1000EA874:'version_error',0x101D5C248:'unsupported',0x1000EABE4:'unsupported',
       0x102226250:'unsupported',0x102226290:'unsupported',0x101B89938:'unsupported'}
STATIC_PAGES={0x102FC2000:65536,0x102D55000:69632,0x1014DF000:73728}

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,_=arm64_image(args.binary.read_bytes())
    source_hash=hashlib.sha256(image).hexdigest()
    if source_hash!='2cb4de127b50d54c6c5666acfef2263fae16d11a3361ad39353ecf8a66a4d9c7':
        raise ValueError('Unrecognized 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('Unrecognized constructor input')
        bodies[address]=data
    functions={address:bodies[address] for address,_,_ in RANGES[:-1]}
    compiler=ProgramCompiler(HOSTS,returning_hosts={0x1027D36E4},memory_pages=16,static_pages=STATIC_PAGES)
    wat=compiler.compile(functions,ENTRY).replace('(import "env"','(import "deser_host"')
    out=ROOT/'build/deserialiser';out.mkdir(parents=True,exist_ok=True)
    (out/'factorio_deser.wat').write_text(wat)
    run(ROOT/'build/wabt/wat2wasm',out/'factorio_deser.wat','-o',out/'factorio_deser.wasm')
    run(ROOT/'build/wabt/wasm2c',out/'factorio_deser.wasm','-o',out/'factorio_deser.c','--module-name=factorio_deser')
    assembly=['.text','.p2align 2','.globl _factorio_deser_original']
    for address,data in functions.items():
        if address==ENTRY: assembly.append('_factorio_deser_original:')
        for i,(word,) in enumerate(struct.iter_unpack('<I',data)):
            pc=address+i*4;assembly.append(f'Ldeser_{pc:x}:')
            if word & 0xfc000000 == 0x94000000:
                target=pc+((word&0x3ffffff)-(0x4000000 if word&0x2000000 else 0))*4
                assembly.append('bl '+(f'Ldeser_{target:x}' if target in functions else '_native_deser_'+HOSTS[target]))
            elif pc==0x10244A50C: assembly.append('adrp x9, _native_deser_vtable@PAGE')
            elif pc==0x10244A510: assembly.append('add x9, x9, _native_deser_vtable@PAGEOFF')
            else: assembly.append(f'.long 0x{word:08x}')
    assembly += ['.section __DATA,__const','.p2align 3','.globl _native_deser_vtable',
                 '_native_deser_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_deser.c',ROOT/'runtime/deserialiser_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/deserialiser_equivalence.c','-o',out/'verify_deserialiser')
    run(out/'verify_deserialiser')
    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_deser_ios.a',*objects)
    manifest={'arm64_sha256':source_hash,'ranges':[{'address':hex(a),'bytes':n,'sha256':h} for a,n,h in RANGES],
              'instructions':compiler.instructions,'calls':compiler.calls,'memory_bytes':1048576,
              'scope':'original constructor initialization and version decisions; host reads/allocation/bzero; no world loading or unwinding',
              'exceptions':'abort status 7 at cxa_allocate_exception; original exception objects are not constructed',
              'static_pages':{hex(a):b for a,b in STATIC_PAGES.items()}}
    (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,
        'normal_boundary_cases':3,'truncated_headers':9,'version_errors':8,'allocation_failures':12,
        'instructions':len(compiler.instructions),'owned_target_allocations':12,'owned_target_heap_bytes':10712,
        'scope':'constructor initialization and explicit abort paths; not world loading or C++ unwinding'},indent=2)+'\n')
    print('PASS: original MapDeserialiser initialization graph, desktop comparisons and iOS library')

if __name__=='__main__':main()