#!/usr/bin/env python3
"""Recover bounded radio parser states, key comparisons and substitution branches."""
import json
from pathlib import Path
from analyze_deferred_auth import partition_function, trace_slice
from analyze_game_cvars import vm_string
from analyze_qvm import decode, write_json
from analyze_vm_extensions import verified

BASE=0x40000000


def compare(ins,left,right,bound=99999):
    a,b=0x10000,0x20000
    memory={a+i:v for i,v in enumerate(left+b'\0')}
    memory.update({b+i:v for i,v in enumerate(right+b'\0')})
    return trace_slice(ins,69387,69525,{}, {48:a,52:b,56:bound},
                       byte_memory=memory,limit=4096,capture_stack=True,halt_on_leave=True)


def main():
    q=json.loads(Path('docs/research/qvm/qagame.json').read_text())
    data=verified(Path(q['source']['member']['output']).read_bytes(),q['sha256']);h,ins=decode(data)
    functions={str(e):partition_function(ins,e,data,h) for e in [69386,185212,186435]}
    keys=[]
    for i in range(185212,186435):
        if ins[i]['opcode']=='CONST' and ins[i]['operand']==69386 and ins[i+1]['opcode']=='CALL':
            literal=ins[i-5]
            if literal['opcode']!='CONST' or ins[i-4]['opcode']!='ARG' or ins[i-4]['operand']!=12:
                raise ValueError('key comparison pattern changed')
            keys.append({'key':vm_string(data,h,literal['operand']),'address':literal['operand'],
                         'compare_call':i+1,'bound':ins[i-3]['operand'],'window':ins[i-11:i+8]})
    comparisons=[]
    for k in keys:
        key=k['key'].encode('ascii')
        for candidate in [key,key.upper(),key[:-1],key+b'x']:
            t=compare(ins,candidate,key)
            comparisons.append({'key':k['key'],'candidate':candidate.decode('ascii'),'trace':t})
    # Per-character slices: no copy/append/lookup callee is entered.
    chars=[]
    for marker,key,ch in [(0,'',ord('x')),(0,'',36),(36,'foo',35),(35,'',ord('a')),
                           (36,'h',ord('P')),(36,'foo',ord('1')),(36,'foo',32),
                           (36,'foo',255),(0,'',0)]:
        bm={BASE+32:ch,BASE+182:marker}
        bm.update({BASE+183+i:v for i,v in enumerate(key.encode()+b'\0')})
        t=trace_slice(ins,186481,186676,{}, {28:0,336:len(key),376:0x50000,380:0x60000},byte_memory=bm)
        chars.append({'pending_marker':marker,'pending_key':key,'byte':ch,'trace':t})
    lookup_results=[]
    for result in [0,0x70000]:
        t=trace_slice(ins,186609,186676,{}, {340:result,380:0x60000},byte_memory={BASE+182:36})
        lookup_results.append({'supplied_lookup_return':result,'trace':t})
    tails=[]
    for marker in [0,35,36]:
        t=trace_slice(ins,186693,186719,{}, {380:0x60000},byte_memory={BASE+182:marker})
        tails.append({'pending_marker':marker,'trace':t})
    lengths=[]
    for length in [149,150,151]:
        t=trace_slice(ins,186449,186719,{}, {344:length,380:0x60000},
                       byte_memory={BASE+182:0},limit=256)
        lengths.append({'supplied_source_length':length,'trace':t})
    entity,record=0x10000,0x20000
    health=[]
    for value in [-1,0,24,25,49,50,74,75,100]:
        for state,override in [(0,0),(3,0),(4,0),(0,1)]:
            t=trace_slice(ins,185237,186431,{entity+684:value,entity+520:record,
                        record+4:state,record+4324:override},{1144:entity})
            health.append({'raw_health':value,'raw_state':state,'override':override,'trace':t})
    hp=[]
    for state,override in [(0,0),(3,0),(4,0),(0,1)]:
        t=trace_slice(ins,185350,186431,{entity+684:42,entity+520:record,record+4:state,
                        record+4324:override},{1144:entity})
        hp.append({'raw_state':state,'override':override,'trace':t})
    side_effects=[]
    for pointer,now in [(0,1000),(0x30000,1000),(0x30000,2147483640)]:
        t=trace_slice(ins,185987,186031,{entity+520:record,record+6168:pointer,34992532:now},
                      {1144:entity},byte_memory={},capture_stack=True,halt_on_leave=True)
        side_effects.append({'object_pointer':pointer,'now':now,'trace':t})
    addresses=[77862,77857,77850,77847,78120,78109,78095,78087,78079,78074,78068,
               78066,78037,78014,78009,78000,77991,77972,77959,77936]
    addresses.extend(k['address'] for k in keys)
    strings={str(a):vm_string(data,h,a) for a in addresses}
    out={'source':q['source'],'sha256':q['sha256'],'functions':functions,'keys':keys,
         'comparison_traces':comparisons,'character_traces':chars,'lookup_result_traces':lookup_results,
         'pending_tail_traces':tails,'source_length_traces':lengths,'health_traces':health,'hp_traces':hp,
         'momo_traces':side_effects,'strings':strings,
         'limitations':'Static partitions and bounded synthetic instruction slices. No original process or full '
         'QVM executes. CALLs return zero placeholders; source lengths and isolated lookup returns are supplied '
         'inputs, not executed callees. No composed parser output or formatting/copy/append execution. '
         'Timer, weapon, location, leader and crosshair bodies are retained but not fully evaluated; field '
         'identities, ownership, marker lifetime and original runtime/platform behavior remain open.'}
    write_json(Path('docs/research/chat-expansion.json'),out)
    print(f'Recovered {len(keys)} substitution keys, parser state and health/momo slices')


if __name__=='__main__':main()