#!/usr/bin/env python3
"""Synthetic unrolled/recursive send and supplied-tree symbol round trips."""
import copy
import gzip
import itertools
import json
from pathlib import Path
from analyze_area_queries import site
from analyze_bsp_references import global_address,trace_fields
from analyze_collision_composition import evaluate,get_word,put_word
from analyze_engine_control import control_flow
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions,verified

DATA,CURSOR,TREE,TABLE,OUT=0x10000,0x20000,0x30000,0x50000,0x60000


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


def send(g,m,parent,child,buffer,stack=0x80000,level=1):
    rows=g['instructions'];start=site(g,'movq','%rdi, 0x20(%rsp)');end=site(g,'movq','0x28(%rsp), %rbx');idx=next(i for i,x in enumerate(rows) if x['opcode']=='callq');call=rows[idx]['address']
    regs={'%rdi':parent,'%rsi':child,'%rdx':buffer,'%r12':buffer,'%rsp':stack};trace=[];frames=1;maximum=level
    while True:
        t=run(g,start,[call,end],m,regs);trace+=t['trace'];m=t['memory'];regs=t['registers']
        if t['stop']==end:break
        child_result=send(g,m,regs['rdi'],regs['rsi'],regs['rdx'],stack+0x1000,level+1)
        m=child_result['memory'];trace+=child_result['trace'];frames+=child_result['frames'];maximum=max(maximum,child_result['maximum_depth'])
        # Explicit restoration of all caller registers; no implicit ABI or native CALL.
        start=rows[idx+1]['address']
    return {'memory':m,'trace':trace,'frames':frames,'maximum_depth':maximum}


def transmit(graphs,m,symbol,bloc):
    g=copy.deepcopy(graphs['Huff_offsetTransmit']);rows=g['instructions']
    for x in rows:
        if x['opcode']=='pushq':x['opcode']='nop';x['operands']=''
    idx=next(i for i,x in enumerate(rows) if x['opcode']=='callq');end=next(x['address'] for x in rows if x['opcode']=='popq')
    t=run(g,rows[0]['address'],[rows[idx]['address'],end],m,{'%rdi':TABLE,'%rsi':symbol,'%rdx':DATA,'%rcx':CURSOR});trace=t['trace'];m=t['memory'];maximum=0
    if t['stop']!=end:
        regs=t['registers'];c=send(graphs['send'],m,regs['rdi'],regs['rsi'],regs['rdx']);m=c['memory'];trace+=c['trace'];maximum=c['maximum_depth']
        t=run(g,rows[idx+1]['address'],[end],m,regs);trace+=t['trace'];m=t['memory']
    return m,trace,maximum


def comb(m,bits):
    # Chain toward the selected leaf, with a distinct leaf sibling on every level.
    nodes=[]
    for i,bit in enumerate(bits):
        addr=TREE+128*i;nxt=addr+128;sibling=addr+64;left,right=(sibling,nxt) if bit else (nxt,sibling)
        for node,l,r,parent,symbol in [(addr,left,right,addr-128 if i else 0,257),(sibling,0,0,addr,80+i)]:
            put_word(m,node,l,8);put_word(m,node+8,r,8);put_word(m,node+16,parent,8);put_word(m,node+52,symbol);nodes.append([node,l,r,parent,symbol])
    leaf=TREE+128*len(bits);put_word(m,leaf,0,8);put_word(m,leaf+8,0,8);put_word(m,leaf+16,leaf-128 if bits else 0,8);put_word(m,leaf+52,42);nodes.append([leaf,0,0,leaf-128 if bits else 0,42]);put_word(m,TABLE+32+8*42,leaf,8)
    return leaf,nodes


def buffer_memory(start,sentinel):
    m={DATA+i:sentinel for i in range(-1,17)};m[DATA-1]=m[DATA+16]=0x7b;put_word(m,CURSOR,start);return m


def receive(g,m):
    stops=[x['address'] for x in g['instructions'] if x['opcode']=='retq']
    return run(g,g['instructions'][0]['address'],stops,m,{'%rdi':TREE,'%rsi':OUT,'%rdx':DATA,'%rcx':CURSOR})


def main():
    out={'native':{},'limitations':'Paired Linux ELF64 synthetic instruction compositions only. Trees/tables/virtual buffers and disjoint stack frames supplied. send prologue/epilogue excluded; wrapper PUSH omitted; all callers restored explicitly across stopped CALL boundaries. No original process/native callee or actual tree initialization/codebook/message packet executes. Cycles, invalid parent/child links, capacity/truncation/overflow, ABI, real codecs and all-platform/runtime conformance remain open.'}
    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());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());graphs={n:control_flow(fs[n]) for n in ['send','Huff_offsetTransmit','Huff_offsetReceive']}
        bloc=global_address(next(x for x in graphs['send']['instructions'] if '(%rip)' in x['operands']));probes=[];roundtrips=[]
        for depth,pattern,sentinel,start in itertools.product(list(range(11))+[16,17,24,32],range(4),[0,255],[0,7,8]):
            bits=[0 if pattern==0 else 1 if pattern==1 else i%2 if pattern==2 else (i+1)%2 for i in range(depth)];m=buffer_memory(start,sentinel);leaf,nodes=comb(m,bits);put_word(m,bloc,start)
            if depth:t=send(graphs['send'],m,leaf-128,leaf,DATA)
            else:t={'memory':m,'trace':[],'frames':0,'maximum_depth':0}
            m=t['memory'];probes.append({'depth':depth,'path_bits':bits,'sentinel':sentinel,'initial_bit':start,'encoded_bytes':[m[DATA+i] for i in range(16)],'final_global':get_word(m,bloc),'canaries':[m[DATA-1],m[DATA+16]],'frames':t['frames'],'maximum_depth':t['maximum_depth'],**trace_fields(t,False)})
        for depth,pattern,start in itertools.product([0,1,2,8,9,10,17,32],range(4),[0,7,8]):
            bits=[0 if pattern==0 else 1 if pattern==1 else i%2 if pattern==2 else (i+1)%2 for i in range(depth)];m=buffer_memory(start,0);leaf,nodes=comb(m,bits);m,trace,maximum=transmit(graphs,m,42,bloc);end=get_word(m,CURSOR);encoded=[m[DATA+i] for i in range(16)];put_word(m,CURSOR,start)
            t=receive(graphs['Huff_offsetReceive'],m);m=t['memory'];trace+=t['trace']
            roundtrips.append({'depth':depth,'path_bits':bits,'initial_bit':start,'nodes':nodes,'encoded_bytes':encoded,'transmit_cursor':end,'received_symbol':get_word(m,OUT),'receive_cursor':get_word(m,CURSOR),'final_global':get_word(m,bloc),'canaries':[m[DATA-1],m[DATA+16]],'maximum_depth':maximum,**trace_fields({'trace':trace},False)})
        sequences=[]
        for sequence,start in itertools.product([[],[40],list(range(40,48))*2,[47]*16,[40,47]*8],[0,7,8]):
            m=buffer_memory(start,0);nodes=[]
            for i in range(15):
                addr=TREE+64*i;internal=i<7;left=TREE+64*(2*i+1) if internal else 0;right=TREE+64*(2*i+2) if internal else 0;parent=TREE+64*((i-1)//2) if i else 0;symbol=257 if internal else 40+i-7
                put_word(m,addr,left,8);put_word(m,addr+8,right,8);put_word(m,addr+16,parent,8);put_word(m,addr+52,symbol);nodes.append([addr,left,right,parent,symbol])
                if not internal:put_word(m,TABLE+32+8*symbol,addr,8)
            trace=[];cursors=[]
            for symbol in sequence:m,part,_=transmit(graphs,m,symbol,bloc);trace+=part;cursors.append(get_word(m,CURSOR))
            encoded=[m[DATA+i] for i in range(16)];put_word(m,CURSOR,start);decoded=[];readcursors=[]
            for _ in sequence:
                t=receive(graphs['Huff_offsetReceive'],m);m=t['memory'];trace+=t['trace'];decoded.append(get_word(m,OUT));readcursors.append(get_word(m,CURSOR))
            sequences.append({'symbols':sequence,'initial_bit':start,'nodes':nodes,'encoded_bytes':encoded,'decoded_symbols':decoded,'transmit_cursors':cursors,'receive_cursors':readcursors,'canaries':[m[DATA-1],m[DATA+16]],**trace_fields({'trace':trace},False)})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'send_probes':probes,'roundtrips':roundtrips,'sequences':sequences}
    write_json(Path('docs/research/huffman-encoding.json'),out);print('Recovered 720 send compositions, 192 symbol round trips and 30 shared-tree sequences')


if __name__=='__main__':main()