#!/usr/bin/env python3
"""Keyed delta masks/XOR and modeled-tree compressed compositions."""
import gzip
import hashlib
import itertools
import json
import re
import struct
from pathlib import Path
from analyze_bsp_loading import read_address
from analyze_native import elf
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, install_tree
from analyze_message_numeric import frame_graph
from analyze_message_widths import reader, run, writer
from analyze_message_wrappers import WRAPPER_STACK
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions, verified


INTEGER_PAIRS=[(0,0),(0,1),(1,0),(0x1234,0x1234),(0x1234,0xffffffff),(0xffffffff,0)]
FLOAT_PAIRS=[(0,0),(0,0x80000000),(0x80000000,0),(0x3f800000,0x3f800000),(0x3f800000,0x3fc00000),(0xbf800000,0x3f800000),(0x3f800000,0),(0x7f7fffff,0x7f7fffff)]


def write_delta(graphs,m,kind,baseline,value,width,key):
    name='MSG_WriteDeltaKeyFloat' if kind=='float' else 'MSG_WriteDeltaKey';g,omitted=frame_graph(graphs[name]);rows=g['instructions'];calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'};tails={x['address'] for x in rows if x['opcode']=='jmp' and '<MSG_WriteBits>' in x['operands']};ends={x['address'] for x in rows if x['opcode']=='retq'}
    regs={'%rdi':MSG,'%rsi':key,'%rdx':baseline,'%rcx':value,'%r8':width&0xffffffff,'%xmm0':baseline,'%xmm1':value,'%rsp':WRAPPER_STACK};pc=rows[0]['address'];requests=[];trace=[]
    while True:
        t=run(g,pc,list(calls)+list(tails)+list(ends),m,regs);m=t['memory'];regs=t['registers'];trace+=t['trace']
        if t['stop'] in ends:break
        request=[regs[k] for k in ['rdi','rsi','rdx']];requests.append(request)
        if request[0]!=MSG:raise ValueError('delta writer pointer')
        word_width=request[2];signed_width=word_width if word_width<2**31 else word_width-2**32
        m,_,error,child=writer(graphs,m,signed_width,request[1]);trace+=child
        if error:raise ValueError('delta writer fixture error')
        if t['stop'] in tails:break
        pc=rows[calls[t['stop']]+1]['address']
        if len(requests)>2:raise ValueError('delta writer bound')
    return m,requests,trace,omitted


def read_delta(graphs,m,kind,baseline,width,key,supplied=None):
    name='MSG_ReadDeltaKeyFloat' if kind=='float' else 'MSG_ReadDeltaKey';g,omitted=frame_graph(graphs[name]);rows=g['instructions'];calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'};tails={x['address'] for x in rows if x['opcode']=='jmp' and '<MSG_ReadBits>' in x['operands']};ends={x['address'] for x in rows if x['opcode']=='retq'}
    regs={'%rdi':MSG,'%rsi':key,'%rdx':baseline,'%rcx':width&0xffffffff,'%xmm0':baseline,'%rsp':WRAPPER_STACK};pc=rows[0]['address'];requests=[];trace=[]
    while True:
        t=run(g,pc,list(calls)+list(tails)+list(ends),m,regs);m=t['memory'];regs=t['registers'];trace+=t['trace']
        if t['stop'] in ends:break
        request=[regs[k] for k in ['rdi','rsi']];requests.append(request)
        if request[0]!=MSG:raise ValueError('delta reader pointer')
        word_width=request[1];signed_width=word_width if word_width<2**31 else word_width-2**32
        if supplied is None:
            m,value,error,child=reader(graphs,m,signed_width);trace+=child
            if error:raise ValueError('delta reader fixture error')
        else:value=supplied[len(requests)-1]
        regs['rax']=value
        if t['stop'] in tails:break
        pc=rows[calls[t['stop']]+1]['address']
        if len(requests)>2:raise ValueError('delta reader bound')
    return m,regs['xmm0' if kind=='float' else 'rax']&0xffffffff,requests,trace,omitted


def main():
    cp=Path('docs/research/huffman-codebook.json');ip=Path('docs/research/huffman-initialization.json');code=json.loads(cp.read_text());initial=json.loads(ip.read_text())
    out={'native':{},'codebook_report_sha256':hashlib.sha256(cp.read_bytes()).hexdigest(),'initialization_report_sha256':hashlib.sha256(ip.read_bytes()).hexdigest(),'limitations':'Synthetic Linux ELF64 instruction compositions only; original processes/native callees never execute. Compressed tree is MODEL-DERIVED, not complete native training/runtime capture. Save/restore/PUSH/POP omitted; disjoint virtual frames and caller-register restoration explicit. Float writer comparisons are finite only; signed zeros compare equal. Separate supplied-reader flag/return probes do not compose actual bits or cursor effects. Ample zero-initialized buffers, sampled integer widths and pairs only. Integer mask table has 32 symbol-backed words; sampled width32 reads a file-backed adjacent entityStateFields word, not a valid mask entry. No legal-index/overflow/truncation/usercmd/entity/player/snapshot/real-packet/all-platform conformance.'}
    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']);meta,symbols=elf(data);a=r['artifacts']['disassembly.txt'];fs=native_functions(verified(gzip.open(a['path'],'rb').read(),a['uncompressed_sha256']).decode());names=['MSG_WriteDeltaKey','MSG_ReadDeltaKey','MSG_WriteDeltaKeyFloat','MSG_ReadDeltaKeyFloat','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']!=r['source'] or code['native'][label]['disassembly_sha256']!=a['uncompressed_sha256']:raise ValueError('codebook provenance')
        if initial['native'][label]['source']!=r['source'] or initial['native'][label]['disassembly_sha256']!=a['uncompressed_sha256']:raise ValueError('initialization provenance')
        mask_instruction=next(x for x in graphs['MSG_ReadDeltaKey']['instructions'] if x['opcode']=='andl' and '(,%rax,4)' in x['operands'])
        mask_address=int(re.search(r'(0x[0-9a-f]+)\(',mask_instruction['operands'])[1],16);raw,offset=read_address(data,meta['sections'],mask_address,132)
        mask_symbol=next(x for x in symbols if x['name']=='kbitmask' and x['address']==mask_address);adjacent=next(x for x in symbols if x['name']=='entityStateFields' and x['address']==mask_address+128)
        if mask_symbol['bytes']!=128:raise ValueError('mask symbol size')
        mask_words=list(struct.unpack('<33I',raw));mask_report={'address':mask_address,'file_offset':offset,'symbol':mask_symbol,'table_words':mask_words[:32],'table_sha256':hashlib.sha256(raw[:128]).hexdigest(),'adjacent_word':mask_words[32],'adjacent_symbol':adjacent,'adjacent_bytes_hex':raw[128:].hex()}
        template={mask_address+i:v for i,v in enumerate(raw)};base=initial['native'][label]['entry_request']['huff_init_pointer'];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']))
        streams=[]
        fixtures=[('integer',w,b,v,0,key) for w,(b,v),key in itertools.product([1,7,8,9,16,24,31,32],[(0,0),(0,1),(0x1234,0x1234),(0x1234,0xffffffff)],[0,0x80000100,0xffffffff])]+[('float',32,b,v,bit,key) for (b,v),bit,key in itertools.product(FLOAT_PAIRS,[0,7],[0,0x80000100,0xffffffff])]
        for kind,width,baseline,value,bit,key in fixtures:
            m=dict(template)
            for i in range(256):m[DATA+i]=0
            m[DATA-1]=m[DATA+256]=0x7b
            for offset,v in [(4,0),(8,0),(16,DATA),(24,256),(28,bit//8+1),(32,0),(36,bit)]:put_word(m,MSG+offset,v,8 if offset==16 else 4)
            put_word(m,old,0);put_word(m,over,0)
            m,wrequests,trace,womitted=write_delta(graphs,m,kind,baseline,value,width,key);write_bit=get_word(m,MSG+36);size=get_word(m,MSG+28);payload=bytes(m[DATA+i] for i in range(size)).hex();account=get_word(m,old);counter=get_word(m,over)
            put_word(m,MSG+36,bit);put_word(m,MSG+32,0);m,result,rrequests,child,romitted=read_delta(graphs,m,kind,baseline,width,key);trace+=child
            streams.append({'kind':kind,'key_word':key,'width':width,'baseline_word':baseline,'value_word':value,'initial_bit':bit,'write_requests':wrequests,'read_requests':rrequests,'payload_hex':payload,'write_bit':write_bit,'write_size':size,'read_bit':get_word(m,MSG+36),'read_count':get_word(m,MSG+32),'result_word':result,'oldsize_word':account,'range_counter':counter,'overflow_flag':get_word(m,MSG+4),'canaries':[m[DATA-1],m[DATA+256]],'writer_omitted':womitted,'reader_omitted':romitted,**trace_fields({'trace':trace},False)})
        probes=[]
        for kind,flag,key in itertools.product(['integer','float'],[0,1,2,0xffffffff],[0,0x80000100,0xffffffff]):
            baseline=0x1234;payload=0x8196a5ff
            _,result,requests,trace,omitted=read_delta(graphs,template,kind,baseline,13,key,[flag,payload]);probes.append({'kind':kind,'key_word':key,'supplied_flag_word':flag,'baseline_word':baseline,'supplied_payload_word':payload,'width':13,'result_word':result,'requests':requests,'omitted':omitted,**trace_fields({'trace':trace},False)})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'mask_table':mask_report,'compressed_streams':streams,'supplied_reader_probes':probes}
    write_json(Path('docs/research/message-delta-key.json'),out)
    print('Recovered keyed delta mask/XOR boundaries and modeled-tree streams')


if __name__=='__main__':main()