#!/usr/bin/env python3
"""Bounded native message initialization, OOB and mixed-width compositions."""
import gzip
import hashlib
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 get_word, put_word
from analyze_engine_control import control_flow
from analyze_message_byte_codec import MSG, DATA, STACK, install_tree, offset_transmit, run
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions, verified


def writer(graphs,m,width,value):
    g=graphs['MSG_WriteBits'];rows=g['instructions']
    start=next(x['address'] for x in rows if x['opcode']=='addl' and '(%rip)' in x['operands'])
    end=next(x['address'] for x in rows if x['opcode']=='movq' and x['operands'].split('#')[0].strip()=='0x8(%rsp), %rbx')
    # Stop before the unsupported-OOB-width epilogue restores saved registers.
    bad=[x['address'] for x in rows if x['opcode']=='movl' and x['operands'].split('#')[0].strip()=='%r14d, %edx'][-1]
    calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'}
    regs={'%rdi':MSG,'%rbp':MSG,'%rbx':value,'%rdx':width&0xffffffff,'%rsp':STACK};trace=[];requests=[];error=None
    while True:
        t=run(g,start,list(calls)+[end,bad],m,regs);m=t['memory'];regs=t['registers'];trace+=t['trace']
        if t['stop']==end:break
        if t['stop']==bad:error='oob_width';break
        i=calls[t['stop']];target=rows[i]['operands']
        if '<Com_Error>' in target:error='width';break
        if '<Huff_offsetTransmit>' in target:
            requests.append(['symbol',regs['rsi']]);m,child=offset_transmit(graphs['Huff_offsetTransmit'],graphs['send'],m,regs);trace+=child
        elif '<Huff_putBit>' in target:
            requests.append(['raw',regs['rdi']]);h=graphs['Huff_putBit'];child=run(h,h['instructions'][0]['address'],[x['address'] for x in h['instructions'] if x['opcode']=='retq'],m,regs);m=child['memory'];regs['rax']=child['registers']['rax'];trace+=child['trace']
        else:raise ValueError('unexpected writer helper '+target)
        start=rows[i+1]['address']
    return m,requests,error,trace


def reader(graphs,m,width):
    g=graphs['MSG_ReadBits'];rows=g['instructions'];start=site(g,'testl','%esi, %esi');end=next(x['address'] for x in rows if x['opcode']=='movl' and x['operands'].split('#')[0].strip()=='%ebp, %eax');calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'}
    regs={'%r12':MSG,'%r15':width&0xffffffff,'%rsi':width&0xffffffff,'%rsp':STACK};trace=[];error=None
    while True:
        t=run(g,start,list(calls)+[end],m,regs);m=t['memory'];regs=t['registers'];trace+=t['trace']
        if t['stop']==end:break
        i=calls[t['stop']];target=rows[i]['operands']
        if '<Com_Error>' in target:error='width';break
        name='Huff_getBit' if '<Huff_getBit>' in target else 'Huff_offsetReceive' if '<Huff_offsetReceive>' in target else None
        if name is None:raise ValueError('unexpected reader helper '+target)
        h=graphs[name];child=run(h,h['instructions'][0]['address'],[x['address'] for x in h['instructions'] if x['opcode']=='retq'],m,regs);m=child['memory'];trace+=child['trace']
        if name=='Huff_getBit':regs['rax']=child['registers']['rax']
        start=rows[i+1]['address']
    return m,None if error else regs['rbp']&0xffffffff,error,trace


def main():
    cp=Path('docs/research/huffman-codebook.json');code=json.loads(cp.read_text());init=json.loads(Path('docs/research/huffman-initialization.json').read_text())
    out={'native':{},'codebook_report_sha256':hashlib.sha256(cp.read_bytes()).hexdigest(),'limitations':'Synthetic Linux ELF64 instruction compositions only; original processes and native callees never execute. Warm initialization omits saves/restores; cold initialization stops at first Huff_Init. Explicit caller-register restoration and supplied virtual memory. Compressed tree is MODEL-DERIVED, not full native training/runtime capture. Oversized buffers, selected bit/value states only; no capacity/truncation/lifecycle/packet/all-platform conformance.'}
    for label in ['linux--Quake3-UrT-Ded.x86_64','linux--Quake3-UrT.x86_64']:
        report=json.loads(Path(f'docs/research/native/{label}.json').read_text());verified((Path('.cache/UrbanTerror43')/report['source']['path']).read_bytes(),report['source']['sha256']);a=report['artifacts']['disassembly.txt'];fs=native_functions(verified(gzip.open(a['path'],'rb').read(),a['uncompressed_sha256']).decode())
        names=['MSG_Init','MSG_InitOOB','MSG_Clear','MSG_WriteBits','MSG_ReadBits','Huff_putBit','Huff_getBit','Huff_offsetTransmit','Huff_offsetReceive','send'];graphs={n:control_flow(fs[n]) for n in names}
        if code['native'][label]['source']!=report['source'] or code['native'][label]['disassembly_sha256']!=a['uncompressed_sha256']:raise ValueError('codebook provenance')
        initialization=[]
        for name,flag,size in itertools.product(['MSG_Init','MSG_InitOOB'],[0,1],[0,64,0xffffffff]):
            g=graphs[name];rows=g['instructions'];load=next(x for x in rows if x['opcode']=='movl' and '<msgInit>' in x['operands']);address=global_address(load);end=next(x['address'] for x in rows if x['opcode']=='movq' and x['operands'].split('#')[0].strip()=='0x8(%rsp), %rbx');call=next(x['address'] for x in rows if '<Huff_Init>' in x['operands'])
            m={MSG+i:0xa5 for i in range(-1,41)};put_word(m,address,flag)
            t=run(g,load['address'],[end,call],m,{'%r13':MSG,'%r15':DATA,'%rdx':size,'%rsp':STACK})
            initialization.append({'function':name,'initial_flag':flag,'size_word':size,'outcome':'cold_huff_init' if t['stop']==call else 'warm_record','final_flag':get_word(t['memory'],address),'record_bytes':[t['memory'][MSG+i] for i in range(40)],'canaries':[t['memory'][MSG-1],t['memory'][MSG+40]],'huff_init_pointer':t['registers']['rdi'] if t['stop']==call else None,**trace_fields(t,False)})
        clears=[]
        for fill in [0,0xa5]:
            g=graphs['MSG_Clear'];m={MSG+i:fill for i in range(-1,41)};t=run(g,g['instructions'][0]['address'],[x['address'] for x in g['instructions'] if x['opcode']=='retq'],m,{'%rdi':MSG});clears.append({'initial_byte':fill,'record_bytes':[t['memory'][MSG+i] for i in range(40)],'canaries':[t['memory'][MSG-1],t['memory'][MSG+40]],**trace_fields(t,False)})
        base=init['native'][label]['entry_request']['huff_init_pointer'];template={};state=code['model_derived_seed_training']['final_state'];install_tree(template,state,base);root=install_tree(template,state,base+0xc830)
        for x in graphs['MSG_ReadBits']['instructions']:
            if x['opcode']=='movq' and '(%rip)' in x['operands']:put_word(template,global_address(x),root,8)
        g=graphs['MSG_WriteBits'];old=global_address(next(x for x in g['instructions'] if x['opcode']=='addl' and '(%rip)' in x['operands']));over=global_address(next(x for x in g['instructions'] if x['opcode']=='addl' and '$0x1,' in x['operands'] and '(%rip)' in x['operands']))
        groups={}
        for mode,widths,values,positions in [(1,[8,16,32,-8,-16,1,7,24,-32],[0,0x8196a5ff],[0,3]),(2,[8,16,32],[0x8196a5ff],[0]),(0,list(range(1,33))+[-1,-7,-8,-9,-15,-16,-17,-24,-31],[0,0x8196a5ff],[0,7])]:
            probes=[]
            for width,value,position in itertools.product(widths,values,positions):
                m=dict(template) if mode==0 else {}
                for i in range(256):m[DATA+i]=0
                m[DATA-1]=m[DATA+256]=0x7b
                initial_size=position if mode else position//8+1;initial_bit=19 if mode else position
                for offset,v in [(4,0),(8,mode),(24,256),(28,initial_size),(32,0),(36,initial_bit)]:put_word(m,MSG+offset,v)
                put_word(m,MSG+16,DATA,8);put_word(m,old,0);put_word(m,over,0)
                m,requests,error,trace=writer(graphs,m,width,value);bit=get_word(m,MSG+36);size=get_word(m,MSG+28);payload=[m[DATA+i] for i in range(size)];counter=get_word(m,over);account=get_word(m,old)
                decoded=None;read_error=None;read_bit=None;read_count=None
                if error is None:
                    put_word(m,MSG+36,initial_bit);put_word(m,MSG+32,position if mode else 0)
                    m,decoded,read_error,child=reader(graphs,m,width);trace+=child;read_bit=get_word(m,MSG+36);read_count=get_word(m,MSG+32)
                probes.append({'width':width,'value_word':value,'mode_word':mode,'position':position,'initial_bit':initial_bit,'initial_size':initial_size,'write_error':error,'requests':requests,'payload_bytes':payload,'write_bit':bit,'write_size':size,'oldsize_word':account,'range_counter':counter,'message_overflow':get_word(m,MSG+4),'decoded_word':decoded,'read_error':read_error,'read_bit':read_bit,'read_count':read_count,'canaries':[m[DATA-1],m[DATA+256]],**trace_fields({'trace':trace},False)})
            groups['compressed' if mode==0 else 'oob' if mode==1 else 'noncanonical_oob']=probes
        out['native'][label]={'source':report['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'initialization':initialization,'clears':clears,**groups}
    write_json(Path('docs/research/message-widths.json'),out)
    print('Recovered initialization/clear, OOB and mixed-width message compositions')


if __name__=='__main__':main()