#!/usr/bin/env python3
"""Bounded reference copies/addressing, area connectivity and portal mutations."""
import gzip
import hashlib
import itertools
import json
from pathlib import Path
import re
from analyze_area_queries import site
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_native import elf
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions,verified


def operand(x):return x['operands'].split('#')[0].strip()
def global_address(x):return int(re.search(r'#\s*0x([0-9a-f]+)',x['operands'])[1],16)
def trace_fields(t,full=False):return {'trace':t['trace'] if full else None,'trace_length':len(t['trace']),'trace_sha256':hashlib.sha256(json.dumps(t['trace'],separators=(',',':')).encode()).hexdigest()}
def run(g,start,stops,m,regs):return evaluate(g,start,stops,m,regs,limit=2048,capture_registers=True)


def errors(g,data,sections):
    result=[];rows=g['instructions']
    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 operand(y).endswith(', %esi'))
        result.append({'call':x['address'],'text':native_string(data,sections,int(operand(fmt).split(',')[0][1:],0))})
    return result


def main():
    out={'native':{},'limitations':'Synthetic supplied-allocation copies and stopped instruction slices. '
         'Original processes, allocations, errors, collision callees and flooding never execute. '
         'Reference array IDs, pointer validity, caller/inline loader equivalence, reset/zeroing, '
         'full flood/portal lifecycle, snapshots and original-runtime/all-platform conformance remain open.'}
    source,output,area_base,matrix,cvar=0x10000,0x30000,0x50000,0x60000,0x70000
    names=['CMod_LoadLeafBrushes','CMod_LoadLeafSurfaces','CM_TraceThroughLeaf','CM_TestInLeaf','CM_AreasConnected','CM_AdjustAreaPortalState','CM_FloodAreaConnections','CM_FloodArea_r']
    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 names}
        references={}
        for name in names[:2]:
            g=graphs[name];rows=g['instructions'];alloc=next(i for i,x in enumerate(rows) if '<Hunk_Alloc>' in x['operands']);start=rows[alloc+1]['address'];end=site(g,'addq','$0x8, %rsp')
            allocation_start=site(g,'leal','0x1(%rbx), %edi') if name.endswith('Brushes') else site(g,'movslq','%ebx, %rdi')
            count_slot=global_address(next(x for x in rows if x['opcode']=='movl' and '(%rip)' in x['operands'] and operand(x).startswith('%ebx,')))
            base_slot=global_address(next(x for x in rows if x['opcode']=='movq' and '(%rip)' in x['operands'] and operand(x).startswith('%rax,')))
            probes=[]
            for count in list(range(10))+[16,17]:
                m={};words=[[-1,0,7,7,2][i%5]&0xffffffff for i in range(count)]
                for i,v in enumerate(words):put_word(m,source+4*i,v)
                for i in range(4*(count+1)+2):m[output-1+i]=0xa5
                allocation=run(g,allocation_start,[rows[alloc]['address']],m,{'%rbx':count})
                t=run(g,start,[end],m,{'%rax':output,'%rbx':count,'%rbp':count,'%r13':0,'%r12':source})
                probes.append({'count':count,'input_words':words,'output_words':[get_word(t['memory'],output+4*i) for i in range(count)],'allocation_bytes':allocation['registers']['rdi'],
                    'count_word':get_word(t['memory'],count_slot),'base_word':get_word(t['memory'],base_slot,8),'canaries':[t['memory'][output-1],t['memory'][output+4*count]],'writes':[w for w in t['writes'] if output<=w['address']<output+4*(count+1)],**trace_fields(t,count<=3)})
            references[name]={'count_slot':count_slot,'base_slot':base_slot,'probes':probes,'errors':errors(g,data,metadata['sections']),
                'direct_CM_LoadMap_calls':[x['address'] for x in control_flow(fs['CM_LoadMap'])['instructions'] if '<'+name+'>' in x['operands']]}
        # Connect leaf fields to the first indexed brush and surface address calculations.
        consumers=[];leaf,brush_base,patch_table=0x80000,0x90000,0xa0000
        g=graphs['CM_TraceThroughLeaf'];rows=g['instructions'];start=next(x['address'] for x in rows if x['opcode']=='movl' and operand(x)=='%ebp, %eax')
        # Two MOVL EBP,EAX sites exist (brush and surface); select the first explicitly.
        brush_stop=next(x['address'] for x in rows if x['opcode']=='cmpl' and operand(x)=='%eax, 0x30(%rbx)')
        brush_base_slot=global_address(next(x for x in rows if x['opcode']=='addq' and '(%rip)' in x['operands'] and operand(x).endswith(', %rbx')))
        brush_refs_slot=references['CMod_LoadLeafBrushes']['base_slot']
        stamp_load=next(x for x in rows if x['opcode']=='movl' and '(%rip)' in x['operands'] and operand(x).endswith(', %eax'))
        for first,index,value in itertools.product([0,2],[0,1],[-1,0,7]):
            m={};put_word(m,leaf+8,first);put_word(m,brush_refs_slot,source,8);put_word(m,source+4*(first+index),value);put_word(m,brush_base_slot,brush_base,8);put_word(m,global_address(stamp_load),41)
            t=run(g,start,[brush_stop],m,{'%rbp':index,'%r12':leaf})
            consumers.append({'kind':'brush','first':first,'ordinal':index,'supplied_id':value,'pointer':t['registers']['rbx'],**trace_fields(t,True)})
        g=graphs['CM_TestInLeaf'];rows=g['instructions'];start=site(g,'movl','%ebp, %eax');end=site(g,'testq','%rbx, %rbx')
        patch_slot=global_address(next(x for x in rows if x['opcode']=='movq' and '(%rip)' in x['operands'] and operand(x).endswith(', %r9')))
        for first,index,value in itertools.product([0,2],[0,1],[-1,0,7]):
            m={};put_word(m,leaf+16,first);put_word(m,references['CMod_LoadLeafSurfaces']['base_slot'],source,8);put_word(m,source+4*(first+index),value);put_word(m,patch_slot,patch_table,8);put_word(m,patch_table+8*value,0xb0000+16*(value+1),8)
            t=run(g,start,[end],m,{'%rbp':index,'%r13':leaf})
            consumers.append({'kind':'surface','first':first,'ordinal':index,'supplied_id':value,'pointer':t['registers']['rbx'],**trace_fields(t,True)})
        # Area predicate: start after frame setup, stop before stack restore or error call.
        g=graphs['CM_AreasConnected'];rows=g['instructions'];start=next(x['address'] for x in rows if x['opcode']=='movq' and '(%rip)' in x['operands']);end=site(g,'movq','0x8(%rsp), %rbx');error=errors(g,data,metadata['sections'])
        noareas_slot=global_address(next(x for x in rows if x['opcode']=='movq' and '(%rip)' in x['operands']));count_slot=global_address(next(x for x in rows if x['opcode']=='movl' and '(%rip)' in x['operands']));area_slot=global_address(next(x for x in rows if x['opcode']=='movq' and '(%rip)' in x['operands'] and operand(x).endswith(', %rsi')))
        connected=[]
        for count,a1,a2,disabled in itertools.product([0,1,3],[-1,0,1,2,3],[-1,0,1,2,3],[0,1,7]):
            m={};put_word(m,noareas_slot,cvar,8);put_word(m,cvar+48,disabled);put_word(m,count_slot,count);put_word(m,area_slot,area_base,8)
            for i,v in enumerate([1,1,2]):put_word(m,area_base+8*i,v);put_word(m,area_base+8*i+4,999)
            t=run(g,start,[end]+[e['call'] for e in error],m,{'%rdi':a1,'%rsi':a2})
            connected.append({'area_count':count,'area1':a1,'area2':a2,'disabled_word':disabled,'error':t['stop']!=end,'result_word':t['registers']['rdx'] if t['stop']==end else None,**trace_fields(t,count==3 and disabled==0 and a1==0)})
        # Portal updates: replace only PUSH frame traffic; stop before POP, errors or flooding.
        g=graphs['CM_AdjustAreaPortalState'];rows=g['instructions'];model=dict(g);model['instructions']=[dict(x) for x in rows];omitted=[]
        for x in model['instructions']:
            if x['opcode']=='pushq':omitted.append(x['address']);x['opcode']='nop';x['operands']=''
        if len(omitted)!=3:raise ValueError('portal frame changed')
        end=next(x['address'] for x in rows if x['opcode']=='popq');flood=next(x['address'] for x in rows if '<CM_FloodArea_r>' in x['operands']);error2=errors(g,data,metadata['sections']);matrix_slot=global_address(next(x for x in rows if x['opcode']=='addq' and '(%rip)' in x['operands']));epoch_slot=global_address(next(x for x in rows if x['opcode']=='addl' and '(%rip)' in x['operands']))
        portal=[]
        for a1,a2,opened,initial in itertools.product([-1,0,1,2,3],[-1,0,1,2,3],[0,1,7,-1],[-1,0,1,2]):
            count=3;m={};put_word(m,count_slot,count);put_word(m,matrix_slot,matrix,8);put_word(m,area_slot,area_base,8);put_word(m,epoch_slot,40)
            for i in range(9):put_word(m,matrix+4*i,0)
            for i in range(3):put_word(m,area_base+8*i,0);put_word(m,area_base+8*i+4,0)
            if 0<=a1<count and 0<=a2<count:
                for i in [a1*count+a2,a2*count+a1]:put_word(m,matrix+4*i,initial)
            t=run(model,model['entry'],[end,flood]+[e['call'] for e in error2],m,{'%rdi':a1,'%rsi':a2,'%rdx':opened})
            outcome='return' if t['stop']==end else 'flood_request' if t['stop']==flood else 'error'
            portal.append({'area1':a1,'area2':a2,'open_word':opened,'initial_counter':initial,'outcome':outcome,'error_text':next((e['text'] for e in error2 if e['call']==t['stop']),None),
                'matrix_words':[get_word(t['memory'],matrix+4*i) for i in range(9)],'matrix_writes':[w['address']-matrix for w in t['writes'] if matrix<=w['address']<matrix+36], 'epoch_word':get_word(t['memory'],epoch_slot),
                'flood_arguments':[t['registers']['rdi'],t['registers']['rsi']] if outcome=='flood_request' else None,**trace_fields(t,a1==0 and a2==1 and opened in [0,1] and initial==1)})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'references':references,'consumer_probes':consumers,
            'area_slots':{'count':count_slot,'base':area_slot,'noareas_cvar_pointer':noareas_slot,'matrix':matrix_slot,'epoch':epoch_slot},'connectivity_errors':error,'connectivity_probes':connected,
            'portal_errors':error2,'portal_omitted_pushes':omitted,'portal_probes':portal}
    write_json(Path('docs/research/bsp-references.json'),out)
    print('Recovered 48 reference copies, 48 consumer-address, 450 connectivity and 800 portal probes')


if __name__=='__main__':main()