#!/usr/bin/env python3
"""Compose bounded native flood instructions and area-bit output slices."""
import gzip
import itertools
import json
from pathlib import Path
from analyze_area_queries import site
from analyze_bsp_references import errors,global_address,operand,trace_fields
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 matrix_for(count,pattern):
    m=[[0]*count for _ in range(count)]
    for a in range(count):
        for b in range(count):
            connected=(a!=b and (pattern=='clique' or pattern=='chain' and abs(a-b)==1 or pattern=='star' and min(a,b)==0 or pattern=='split' and a//3==b//3))
            m[a][b]=2 if connected else -1 if pattern=='negative' and a!=b else 0
    return m


def model(g):
    result=dict(g);result['instructions']=[dict(x) for x in g['instructions']];omitted=[]
    for x in result['instructions']:
        if x['opcode']=='pushq':omitted.append(x['address']);x['opcode']='nop';x['operands']=''
    return result,omitted


def run(g,start,stops,m,regs):return evaluate(g,start,stops,m,regs,limit=50000,capture_registers=True)


def main():
    out={'native':{},'limitations':'Bounded synthetic composition of native instruction graphs, not '
         'original process/callee execution. PUSH traffic is omitted, stack addresses are supplied, '
         'recursive boundaries evaluate a separate synthetic graph and restore the supplied caller '
         'register state; ABI/clobber/frame semantics are not validated. Memset effects are supplied. '
         'Real maps, reset/ownership, malformed/asymmetric states, epoch wrap reachability, snapshots '
         'and original-runtime/all-platform conformance remain open.'}
    areas,matrix,cvar,buffer=0x10000,0x20000,0x30000,0x40000
    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 ['CM_FloodAreaConnections','CM_FloodArea_r','CM_WriteAreaBits','CM_AdjustAreaPortalState']};outer=graphs['CM_FloodAreaConnections'];child=graphs['CM_FloodArea_r'];bits=graphs['CM_WriteAreaBits']
        count_slot=global_address(next(x for x in outer['instructions'] if x['opcode']=='movl' and '(%rip)' in x['operands']))
        epoch_slot=global_address(next(x for x in outer['instructions'] if x['opcode']=='addl' and '(%rip)' in x['operands']))
        area_slot=global_address(next(x for x in outer['instructions'] if x['opcode']=='addq' and '(%rip)' in x['operands'] and operand(x).endswith(', %rdx')))
        matrix_slot=global_address(next(x for x in outer['instructions'] if x['opcode']=='addq' and '(%rip)' in x['operands'] and operand(x).endswith(', %r12')))
        models={};omitted={};ends={};calls={};error_list=errors(child,data,metadata['sections'])
        for name,g in [('CM_FloodAreaConnections',outer),('CM_FloodArea_r',child),('CM_AdjustAreaPortalState',graphs['CM_AdjustAreaPortalState'])]:
            models[name],omitted[name]=model(g);rows=g['instructions'];ends[name]=next(x['address'] for x in rows if x['opcode']=='popq') if name=='CM_AdjustAreaPortalState' else next(x['address'] for x in rows if x['opcode']=='addq' and operand(x).endswith(', %rsp'))
            calls[name]={x['address']:rows[i+1]['address'] for i,x in enumerate(rows) if '<CM_FloodArea_r>' in x['operands']}
        def seed(count,values,epoch=40):
            m={};put_word(m,count_slot,count);put_word(m,epoch_slot,epoch);put_word(m,area_slot,areas,8);put_word(m,matrix_slot,matrix,8)
            for i in range(max(1,count)):put_word(m,areas+8*i,0);put_word(m,areas+8*i+4,0)
            for i,row in enumerate(values):
                for j,v in enumerate(row):put_word(m,matrix+4*(i*count+j),v)
            return m
        def compose(name,m,regs,depth=0):
            if depth>16:raise ValueError('synthetic recursion bound')
            g=models[name];pc=g['entry'];trace=[];requests=[];maximum=depth
            while pc!=ends[name]:
                t=run(g,pc,[ends[name]]+list(calls[name])+[e['call'] for e in error_list],m,regs);m=t['memory'];regs=t['registers'];trace+=t['trace'];pc=t['stop']
                if pc in [e['call'] for e in error_list]:raise ValueError('unexpected flood consistency error')
                if pc in calls[name]:
                    area,component=regs['rdi']&0xffffffff,regs['rsi']&0xffffffff;requests.append([name,area,component,pc])
                    nested=compose('CM_FloodArea_r',m,{'%rdi':area,'%rsi':component,'%rsp':0x100000-(depth+1)*0x1000},depth+1)
                    m=nested['memory'];trace+=nested['trace'];requests+=nested['requests'];maximum=max(maximum,nested['maximum_depth']);pc=calls[name][pc]
                if len(trace)>200000:raise ValueError('synthetic aggregate instruction bound')
            return {'memory':m,'trace':trace,'requests':requests,'maximum_depth':maximum}
        compositions=[]
        for count,pattern in itertools.product([0,1,2,3,4,8,12,16,32],['empty','chain','clique','star','split','negative']):
            values=matrix_for(count,pattern);t=compose('CM_FloodAreaConnections',seed(count,values),{'%rsp':0x100000})
            compositions.append({'count':count,'pattern':pattern,'labels':[get_word(t['memory'],areas+8*i) for i in range(count)],'stamps':[get_word(t['memory'],areas+8*i+4) for i in range(count)],'epoch':get_word(t['memory'],epoch_slot),'requests':t['requests'],'maximum_depth':t['maximum_depth'],**trace_fields(t,count<=2)})
        portal_sequence=[]
        t=compose('CM_FloodAreaConnections',seed(4,matrix_for(4,'empty')),{'%rsp':0x100000});m=t['memory']
        for a1,a2,opened in [(0,1,1),(0,1,7),(0,1,0),(1,2,1),(2,3,-1),(0,1,0),(2,3,0),(0,0,1),(0,0,0)]:
            t=compose('CM_AdjustAreaPortalState',m,{'%rdi':a1,'%rsi':a2,'%rdx':opened,'%rsp':0x100000});m=t['memory']
            portal_sequence.append({'area1':a1,'area2':a2,'open_word':opened,'matrix_words':[get_word(m,matrix+4*i) for i in range(16)],'labels':[get_word(m,areas+8*i) for i in range(4)],'stamps':[get_word(m,areas+8*i+4) for i in range(4)],'epoch':get_word(m,epoch_slot),'requests':t['requests'],'maximum_depth':t['maximum_depth'],**trace_fields(t)})
        # Entry revisit: same epoch/label returns; same epoch/different label requests error.
        revisit=[]
        for epoch,stamp,label_word,requested in itertools.product([40,0xffffffff],[40,0xffffffff,0],[1,2],[1,2]):
            m=seed(1,[[0]],epoch);put_word(m,areas,label_word);put_word(m,areas+4,stamp)
            t=run(models['CM_FloodArea_r'],models['CM_FloodArea_r']['entry'],[ends['CM_FloodArea_r']]+[e['call'] for e in error_list],m,{'%rdi':0,'%rsi':requested,'%rsp':0x100000})
            revisit.append({'epoch':epoch,'stamp_before':stamp,'label_before':label_word,'requested_label':requested,'error':t['stop']!=ends['CM_FloodArea_r'],'label_after':get_word(t['memory'],areas),'stamp_after':get_word(t['memory'],areas+4),**trace_fields(t,True)})
        wrapping=[]
        for epoch in [0,40,0xfffffffe,0xffffffff]:
            for stamp in [0,1,0xffffffff]:
                m=seed(1,[[0]],epoch);put_word(m,areas,77);put_word(m,areas+4,stamp)
                t=compose('CM_FloodAreaConnections',m,{'%rsp':0x100000})
                wrapping.append({'epoch_before':epoch,'stamp_before':stamp,'epoch_after':get_word(t['memory'],epoch_slot),'label_after':get_word(t['memory'],areas),'stamp_after':get_word(t['memory'],areas+4),**trace_fields(t,True)})
        # Output OR loop; all-bits branch stops at memset and receives explicit supplied writes.
        rows=bits['instructions'];start=next(x['address'] for x in rows if x['opcode']=='movl' and '(%rip)' in x['operands']);end=site(bits,'addq','$0x8, %rsp');fill=next(i for i,x in enumerate(rows) if '<memset@plt>' in x['operands']);noareas_slot=global_address(next(x for x in rows if x['opcode']=='movq' and '(%rip)' in x['operands'] and operand(x).endswith(', %rax')))
        outputs=[]
        for count,area,disabled,initial in itertools.product([0,1,7,8,9,16,17],[-1,0],[0,1,7],[0,0xa0]):
            m=seed(count,[[0]*count for _ in range(count)]);put_word(m,noareas_slot,cvar,8);put_word(m,cvar+48,disabled)
            for i in range(max(1,count)):put_word(m,areas+8*i,i%3+1)
            length=(count+7)//8
            for i in range(length+2):m[buffer-1+i]=initial
            t=run(bits,start,[rows[fill]['address'],end],m,{'%rdi':buffer,'%rsi':area});request=None;trace=t['trace']
            if t['stop']!=end:
                regs=t['registers'];request=[regs['rdi'],regs['rsi'],regs['rdx']]
                if request!=[buffer,255,length]:raise ValueError('unexpected memset request')
                for i in range(length):t['memory'][buffer+i]=255
                t=run(bits,rows[fill+1]['address'],[end],t['memory'],regs);trace+=t['trace']
            outputs.append({'count':count,'area':area,'disabled_word':disabled,'initial_byte':initial,'returned_length':t['registers']['rax'],'output_bytes':[t['memory'][buffer+i] for i in range(length)],'canaries':[t['memory'][buffer-1],t['memory'][buffer+length]],'supplied_memset_request':request,**trace_fields({'trace':trace},True)})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'omitted_pushes':omitted,'slots':{'count':count_slot,'area_base':area_slot,'matrix':matrix_slot,'epoch':epoch_slot,'noareas_cvar_pointer':noareas_slot},'consistency_errors':error_list,'compositions':compositions,'portal_sequence':portal_sequence,'revisit_probes':revisit,'epoch_probes':wrapping,'bit_output_probes':outputs}
    write_json(Path('docs/research/area-flooding.json'),out)
    print('Recovered 108 flood compositions, 18 portal-update steps, 48 revisit/24 epoch probes and 168 area-bit outputs')


if __name__=='__main__':main()