#!/usr/bin/env python3
"""Bounded native plane/node/leaf payload conversion after supplied allocation."""
import gzip
import hashlib
import itertools
import json
from pathlib import Path
import re
from analyze_chat_crosshair import bits
from analyze_chat_paths import native_string
from analyze_collision_composition import evaluate,get_word,put_word
from analyze_engine_control import control_flow
from analyze_frame_timing import read_address
from analyze_native import elf
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions,verified

NORMALS=list(itertools.product([-1.,0.,1.],repeat=3))+[(-0.,0.,0.),(.5,-.25,.75),(1.,-2.,3.),(0.,1.,-2.)]
COUNTS=list(range(1,10))+[16,17,31]


def digest(trace):return hashlib.sha256(json.dumps(trace,separators=(',',':')).encode()).hexdigest()


def main():
    out={'native':{},'limitations':'Synthetic native payload conversion after an explicitly supplied '
         'first allocation return; leaf auxiliary-allocation continuation also supplies a return. No allocation, error, original process or native callee executes; '
         'execution stops before epilogues/later allocations. Entry length/division gates '
         'is retained as static evidence, not executed. No actual BSP file is loaded by these probes. '
         'Complete validation/ownership/reset/endian/exceptional float and runtime/platform contracts remain open.'}
    source,output,plane_base=0x10000,0x30000,0x50000
    for label in ['linux--Quake3-UrT-Ded.x86_64','linux--Quake3-UrT.x86_64']:
        r=json.loads(Path(f'docs/research/native/{label}.json').read_text());data=verified((Path('.cache/UrbanTerror43')/r['source']['path']).read_bytes(),r['source']['sha256'])
        a=r['artifacts']['disassembly.txt'];fs=native_functions(verified(gzip.open(a['path'],'rb').read(),a['uncompressed_sha256']).decode());metadata,_=elf(data)
        graphs={n:control_flow(fs[n]) for n in ['CMod_LoadPlanes','CMod_LoadNodes','CMod_LoadLeafs','CM_LoadMap']};results={};routing={}
        caller=graphs['CM_LoadMap']['instructions']
        for name in ['CMod_LoadPlanes','CMod_LoadNodes','CMod_LoadLeafs']:
            i=next(i for i,x in enumerate(caller) if '<'+name+'>' in x['operands'])
            if caller[i-2]['opcode']!='leaq' or caller[i-2]['operands']!='0x10(%rsp), %rdi' or caller[i-1]['opcode']!='addq':raise ValueError('header routing changed')
            offset=int(caller[i-1]['operands'].split(',')[0][1:],0)
            if offset<8 or (offset-8)%8:raise ValueError('unexpected lump descriptor offset')
            t=evaluate(graphs['CM_LoadMap'],caller[i-2]['address'],[caller[i]['address']],{}, {'%rsp':0x90000},capture_registers=True)
            routing[name]={'header_relative_offset':offset,'lump_index':(offset-8)//8,'descriptor_pointer':t['registers']['rdi'],'trace':t['trace'],'call':caller[i]['address']}

        for name in ['CMod_LoadPlanes','CMod_LoadNodes','CMod_LoadLeafs']:
            g=graphs[name]
            rows=g['instructions'];alloc=[i for i,x in enumerate(rows) if '<Hunk_Alloc>' in x['operands']];first=alloc[0];start=rows[first+1]['address']
            end=rows[alloc[1]]['address'] if name=='CMod_LoadLeafs' else next(x['address'] for x in rows[first+1:] if x['opcode']=='popq')
            static_globals=[];seed={}
            for x in rows:
                if '(%rip)' not in x['operands']:continue
                address=int(re.search(r'#\s*0x([0-9a-f]+)',x['operands'])[1],16)
                static_globals.append({'address':address,'instruction':x['address'],'opcode':x['opcode'],'operands':x['operands']})
                if x['opcode']=='movss':
                    raw,_=read_address(data,metadata['sections'],address,4)
                    for i,v in enumerate(raw):seed[address+i]=v
            if name=='CMod_LoadNodes':
                slots={x['address'] for x in static_globals if x['opcode']=='addq'}
                if len(slots)!=1:raise ValueError('node plane-base slot changed')
                for slot in slots:put_word(seed,slot,plane_base,8)
            elif name=='CMod_LoadLeafs':
                slots={x['address'] for x in static_globals if x['opcode']=='cmpl'}
                if len(slots)!=2:raise ValueError('leaf running maxima changed')
                for slot in slots:put_word(seed,slot,0)
            if name=='CMod_LoadPlanes':allocation_start=next(x['address'] for x in rows[:first] if x['opcode']=='leal' and x['operands']=='0xc(%rbp), %ecx')
            elif name=='CMod_LoadLeafs':allocation_start=next(x['address'] for x in rows[:first] if x['opcode']=='leal' and x['operands']=='0x2(%rbp), %r9d')
            else:allocation_start=next(x['address'] for x in rows[:first] if x['opcode']=='movslq' and x['operands']=='%r14d, %rdi')
            probes=[]
            for count in COUNTS:
                m=dict(seed);disk=[];disk_stride={'CMod_LoadPlanes':16,'CMod_LoadNodes':36,'CMod_LoadLeafs':48}[name]
                native_stride={'CMod_LoadPlanes':20,'CMod_LoadNodes':16,'CMod_LoadLeafs':24}[name]
                for i in range(count):
                    if name=='CMod_LoadPlanes':words=[bits(v) for v in NORMALS[i%len(NORMALS)]]+[bits(i-.5)]
                    elif name=='CMod_LoadNodes':words=[[-1,0,1,2,7][i%5],-i-1,i+1]+[0x1100+i*16+j for j in range(6)]
                    else:words=[[-1,0,2,2,7][i%5],[0,3,1,0][i%4]]+[0x2200+i*16+j for j in range(6)]+[100+i,200+i,300+i,400+i]
                    disk.append([v&0xffffffff for v in words])
                    for j,v in enumerate(words):put_word(m,source+disk_stride*i+4*j,v)
                # Fill destination with sentinels, including native padding and spare records.
                for i in range(native_stride*(count+12)+2):m[output-1+i]=0xa5
                allocation=evaluate(g,allocation_start,[rows[first]['address']],m,{'%rbp':count,'%r14':0,'%r13':source},capture_registers=True)
                regs={'%rax':output,'%rbx':source,'%rbp':count,'%r12':count}
                t=evaluate(g,start,[end],m,regs,limit=4096,capture_registers=True)
                records=[]
                for i in range(count):
                    p=output+native_stride*i
                    if name=='CMod_LoadPlanes':records.append({'words':[get_word(t['memory'],p+4*j) for j in range(4)],'type':get_word(t['memory'],p+16,1),'signbits':get_word(t['memory'],p+17,1),'padding':[t['memory'][p+18],t['memory'][p+19]]})
                    elif name=='CMod_LoadNodes':records.append({'plane_pointer':get_word(t['memory'],p,8),'children':[get_word(t['memory'],p+8),get_word(t['memory'],p+12)]})
                    else:records.append([get_word(t['memory'],p+4*j) for j in range(6)])
                portal_allocation=None
                if name=='CMod_LoadLeafs':
                    continuation=dict(t['registers']);continuation['rax']=0x70000
                    portal_allocation=evaluate(g,rows[alloc[1]+1]['address'],[rows[alloc[2]]['address']],t['memory'],continuation,capture_registers=True)
                destination_writes=[w for w in t['writes'] if output<=w['address']<output+native_stride*(count+12)]
                outside=[w for w in t['writes'] if output-1<=w['address']<output or output+native_stride*count<=w['address']<output+native_stride*(count+12)]
                probes.append({'count':count,'allocation_bytes':allocation['registers']['rdi'],'allocation_trace':allocation['trace'],'input_words':disk,'records':records,'destination_writes':destination_writes,'writes_outside_records':outside,
                    'globals':{str(x['address']):get_word(t['memory'],x['address'],8 if x['opcode']=='movq' else 4) for x in static_globals if x['operands'].split('#')[0].strip().endswith('(%rip)') and x['opcode'] in ['movq','movl'] and all(x['address']+j in t['memory'] for j in range(8 if x['opcode']=='movq' else 4))},
                    'portal_allocation_bytes':portal_allocation['registers']['rdi'] if portal_allocation else None,'portal_allocation_trace':portal_allocation['trace'] if portal_allocation else None,'next_allocation_bytes':t['registers']['rdi'] if name=='CMod_LoadLeafs' else None,'trace':t['trace'] if count<=4 else None,'trace_length':len(t['trace']),'trace_sha256':digest(t['trace'])})
            errors=[]
            for i,x in enumerate(rows):
                if '<Com_Error>' not in x['operands']:continue
                fmt=next(y for y in reversed(rows[max(0,i-4):i]) if y['opcode']=='movl' and y['operands'].split('#')[0].strip().endswith(', %esi'))
                address=int(fmt['operands'].split(',')[0][1:],0)
                errors.append({'call':x['address'],'text':native_string(data,metadata['sections'],address)})
            maxima={}
            if name=='CMod_LoadLeafs':
                for role,reg in [('cluster','%edx'),('area','%eax')]:
                    maxima[role]=next(x['address'] for x in static_globals if x['opcode']=='cmpl' and x['operands'].split('#')[0].strip().endswith(', '+reg))
            results[name]={'maxima_slots':maxima,'start_after_supplied_allocation':start,'stop_before_epilogue_or_allocation':end,'global_references':static_globals,'error_requests':errors,'probes':probes}
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'lump_routing':routing,'loaders':results}
    write_json(Path('docs/research/bsp-loading.json'),out)
    print('Recovered 72 supplied-allocation plane/node/leaf conversion probes')


if __name__=='__main__':main()