#!/usr/bin/env python3
"""Bounded byte-addressed collision copies and world-result decisions."""
import gzip
import itertools
import json
import math
from pathlib import Path
import re
import struct
from analyze_chat_crosshair import bits,number
from analyze_engine_control import control_flow
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions,verified


def register(r):
    aliases={'eax':('rax',32),'ax':('rax',16),'dx':('rdx',16),'edi':('rdi',32),'esi':('rsi',32),'edx':('rdx',32),
             'ecx':('rcx',32),'ebp':('rbp',32),'ebx':('rbx',32),'bx':('rbx',16),'bl':('rbx',8),'al':('rax',8),'sil':('rsi',8),'dl':('rdx',8),'cl':('rcx',8),'dil':('rdi',8),'bpl':('rbp',8)}
    r=r.lstrip('%')
    if r in aliases:return aliases[r]
    if re.fullmatch(r'r\d+d',r):return r[:-1],32
    if re.fullmatch(r'r\d+b',r):return r[:-1],8
    if re.fullmatch(r'r(?:ax|di|si|dx|cx|bp|bx|sp|ip|\d+)',r):return r,64
    if re.fullmatch(r'xmm\d+',r):return r,32
    raise ValueError('unsupported register '+r)


def put_word(m,address,value,width=4):
    for i in range(width):m[address+i]=(value>>(8*i))&255


def get_word(m,address,width=4):
    if any(address+i not in m for i in range(width)):raise ValueError('unknown byte-addressed read')
    return sum(m[address+i]<<(8*i) for i in range(width))


def evaluate(graph,start,stops,memory,registers,limit=128,capture_registers=False):
    rows=graph['instructions'];indices={x['address']:i for i,x in enumerate(rows)}
    regs={register(k)[0]:v for k,v in registers.items()};m=dict(memory);equal=None;ge=greater=below=None;signed_le=signed_lt=negative=None;pc=start;trace=[];writes=[]
    def read_reg(name):
        key,size=register(name)
        if key not in regs:raise ValueError('unknown register '+name)
        return regs[key]&((1<<size)-1)
    def address(arg):
        match=re.fullmatch(r'(-?0x[0-9a-f]+)?\((%\w+)?(?:,(%\w+)(?:,([1248]))?)?\)',arg)
        if not match or not (match[2] or match[3]):raise ValueError('unsupported memory operand')
        return ((read_reg(match[2]) if match[2] else 0)+int(match[1] or '0',0)+(read_reg(match[3])*int(match[4] or '1') if match[3] else 0))&((1<<64)-1)
    def read(arg,width):
        if arg.startswith('$'):return int(arg[1:],0)&((1<<(width*8))-1)
        if arg.startswith('%'):return read_reg(arg)
        return get_word(m,address(arg),width)
    def write(arg,val,width):
        val&=(1<<(width*8))-1
        if arg.startswith('%'):
            key,size=register(arg)
            if size in [8,16]:regs[key]=(regs.get(key,0)&~((1<<size)-1))|val
            else:regs[key]=val  # 32-bit writes clear high bits.
        else:
            a=address(arg);put_word(m,a,val,width);writes.append({'address':a,'bytes':width})
    while pc not in stops:
        if pc not in indices or len(trace)>=limit:raise ValueError('composition slice escaped bound')
        x=rows[indices[pc]];trace.append(pc);op=x['opcode'];args=x['operands'].split('#')[0].strip()
        i=indices[pc];pc=rows[i+1]['address'] if i+1<len(rows) else None
        regs['rip']=pc
        if op in ['nop','nopl','nopw']:pass
        elif op in ['movb','movw','movl','movq','movss','movabsq']:
            a,b=args.split(', ');width=8 if op in ['movq','movabsq'] else 1 if op=='movb' else 2 if op=='movw' else 4;write(b,read(a,width),width)
        elif op in ['movzbl','movzwl']:
            a,b=args.split(', ');write(b,read(a,1 if op=='movzbl' else 2),4)
        elif op in ['movsbl','movswl']:
            a,b=args.split(', ');width=1 if op=='movsbl' else 2;v=read(a,width)&((1<<(width*8))-1);sign=1<<(width*8-1)
            write(b,(v^sign)-sign,4)
        elif op=='sbbl':
            if below is None:raise ValueError('unmodeled carry flag')
            a,b=args.split(', ');write(b,read(b,4)-read(a,4)-int(below),4)
            equal=None;signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op in ['cmovael','cmovbl','cmovlel']:
            condition=signed_le if op=='cmovlel' else ge if op=='cmovael' else below
            if condition is None:raise ValueError('unmodeled conditional move flags')
            a,b=args.split(', ')
            if condition:write(b,read(a,4),4)
        elif op=='movaps':
            a,b=args.split(', ')
            if not a.startswith('%xmm') or not b.startswith('%xmm'):raise ValueError('unsupported vector move')
            write(b,read(a,4),4)  # Scalar low lane only.
        elif op=='notl':
            write(args,~read(args,4),4)  # NOT preserves flags.
        elif op=='cltq':
            v=read_reg('%eax');write('%rax',v if v<2**31 else v-2**32,8)
        elif op=='movslq':
            a,b=args.split(', ');v=read(a,4);write(b,v if v<2**31 else v-2**32,8)
        elif op in ['leal','leaq']:
            a,b=args.split(', ');write(b,address(a),4 if op=='leal' else 8)
        elif op in ['testl','testb','testq','cmpl','cmpq','cmpb']:
            width=1 if op in ['testb','cmpb'] else 8 if op in ['testq','cmpq'] else 4
            a,b=args.split(', ');v,w=read(a,width),read(b,width)
            def signed(value):
                sign=1<<(width*8-1)
                return (value^sign)-sign
            equal=v==w if op in ['cmpl','cmpq','cmpb'] else (v&w)==0
            signed_le=signed(w)<=signed(v) if op in ['cmpl','cmpq','cmpb'] else signed(v&w)<=0
            signed_lt=signed(w)<signed(v) if op in ['cmpl','cmpq','cmpb'] else signed(v&w)<0
            negative=signed((w-v)&((1<<(width*8))-1))<0 if op in ['cmpl','cmpq','cmpb'] else signed(v&w)<0
            ge=w>=v if op in ['cmpl','cmpq','cmpb'] else True
            greater=w>v if op in ['cmpl','cmpq','cmpb'] else not equal
            below=w<v if op in ['cmpl','cmpq','cmpb'] else False
        elif op in ['imull','imulq']:
            width=8 if op=='imulq' else 4;parts=args.split(', ')
            if len(parts)==2:a,b=parts;v,w=read(a,width),read(b,width);dest=b
            elif len(parts)==3:a,b,dest=parts;v,w=read(a,width),read(b,width)
            else:raise ValueError('unsupported multiply operands')
            write(dest,v*w,width);equal=None;signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op=='negl':
            write(args,-read(args,4),4);equal=None;signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op in ['sarl','shll','shrl']:
            a,b=args.split(', ') if ', ' in args else ('$0x1',args);amount=read(a,4)&31;v=read(b,4)
            if amount:
                signed_v=v if v<2**31 else v-2**32
                write(b,signed_v>>amount if op=='sarl' else v>>amount if op=='shrl' else v<<amount,4)
                equal=None;signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op=='shlq':
            a,b=args.split(', ');write(b,read(b,8)<<(read(a,8)&63),8)
            equal=None;signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op=='sarq':
            a,b=args.split(', ');v=read(b,8);v=v if v<2**63 else v-2**64
            write(b,v>>(read(a,8)&63),8);equal=None;signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op in ['xorl','xorb','orl','orb','andl','subl','addl','addq','subq']:
            width=8 if op in ['addq','subq'] else 1 if op in ['orb','xorb'] else 4
            a,b=args.split(', ')
            v,w=(0,0) if op in ['xorl','xorb'] and a==b else (read(a,width),read(b,width))
            result=v^w if op in ['xorl','xorb'] else v|w if op in ['orl','orb'] else v&w if op=='andl' else w-v if op in ['subl','subq'] else v+w
            write(b,result,width);equal=(result&((1<<(width*8))-1))==0
            signed_le=signed_lt=negative=None;ge=greater=below=None
        elif op=='xorps':
            a,b=args.split(', ')
            if a!=b:raise ValueError('unsupported SIMD xor')
            write(b,0,4)  # Does not change integer flags.
        elif op=='andps':
            a,b=args.split(', ');write(b,read(a,4)&read(b,4),4)
        elif op=='maxss':
            a,b=args.split(', ');v,w=number(read(a,4)),number(read(b,4))
            write(b,read(b,4) if w>v else read(a,4),4)
        elif op=='sqrtss':
            a,b=args.split(', ');write(b,bits(math.sqrt(number(read(a,4)))),4)
        elif op in ['addss','subss','mulss']:
            a,b=args.split(', ');v,w=number(read(a,4)),number(read(b,4))
            write(b,bits(w+v if op=='addss' else w-v if op=='subss' else w*v),4)
        elif op=='cvttss2si':
            a,b=args.split(', ')
            if register(b)[1]!=32:raise ValueError('unsupported non-int32 scalar conversion')
            v=number(read(a,4));value=int(v)
            if not -2**31<=value<2**31:raise ValueError('unsupported out-of-range scalar conversion')
            write(b,value,4)  # Finite truncation only; exception/MXCSR state unmodeled.
        elif op=='cvtsi2sd':
            a,b=args.split(', ')
            if a.startswith('%') and register(a)[1]!=32:raise ValueError('unsupported non-int32 double conversion')
            v=read(a,4)&0xffffffff;v=v if v<2**31 else v-2**32
            regs[register(b)[0]]=struct.unpack('<Q',struct.pack('<d',float(v)))[0]
        elif op in ['mulsd','cvtsd2ss']:
            def scalar_double(arg):
                if arg.startswith('%') and register(arg)[0] not in regs:raise ValueError('unknown scalar double register')
                word=regs[register(arg)[0]] if arg.startswith('%') else read(arg,8)
                v=struct.unpack('<d',struct.pack('<Q',word&((1<<64)-1)))[0]
                if not math.isfinite(v):raise ValueError('unsupported nonfinite scalar double')
                return v
            a,b=args.split(', ');v=scalar_double(a)
            if op=='mulsd':
                result=scalar_double(b)*v
                if not math.isfinite(result):raise ValueError('unsupported nonfinite scalar double result')
                regs[register(b)[0]]=struct.unpack('<Q',struct.pack('<d',result))[0]
            else:write(b,bits(v),4)
        elif op=='comiss':
            a,b=args.split(', ');v,w=number(read(a,4)),number(read(b,4));equal=v==w;ge=w>=v;greater=w>v;below=w<v;signed_le=signed_lt=negative=None
        elif op in ['sete','setne','setb','setbe','setae']:
            if op in ['sete','setne'] and equal is None:raise ValueError('unmodeled zero flag')
            if op in ['setb','setbe','setae'] and below is None:raise ValueError('unmodeled unsigned flags')
            write(args,int(equal if op=='sete' else not equal if op=='setne' else below if op=='setb' else ge if op=='setae' else below or equal),1)
        elif op in ['je','jne','jae','ja','jb','jbe','jle','jl','jge','jg','js','jns','jmp']:
            if op in ['je','jne'] and equal is None:raise ValueError('unmodeled zero flag')
            if op in ['jle','jl','jge','jg','js','jns'] and signed_le is None:raise ValueError('unmodeled signed flags')
            if op in ['jae','ja','jb','jbe'] and ge is None:raise ValueError('unmodeled unsigned flags')
            taken=op=='jmp' or (equal if op=='je' else not equal if op=='jne' else signed_le if op=='jle' else signed_lt if op=='jl' else not signed_lt if op=='jge' else not signed_le if op=='jg' else negative if op=='js' else not negative if op=='jns' else greater if op=='ja' else below if op=='jb' else below or equal if op=='jbe' else ge)
            if taken:pc=int(args.split()[0],16)
        else:raise ValueError('unsupported composition opcode '+op)
    result={'trace':trace,'stop':pc,'memory':m,'writes':writes}
    if capture_registers:result['registers']={k:v for k,v in regs.items() if k!='rip'}
    return result


def lane_sites(graph):
    rows=graph['instructions'];starts=[x['address'] for x in rows if x['opcode']=='movl' and x['operands']=='0x1030(%rsp), %r10d']
    comparisons=[i for i,x in enumerate(rows) if x['opcode']=='comiss' and x['operands'].startswith('0x44(%rbx),')]
    if len(starts)!=3 or len(comparisons)!=3:raise ValueError('lane count changed')
    return [{'lane':lane,'start':start,'stop':int(rows[i+1]['operands'].split()[0],16),
             'candidate_register':'%r12' if lane==1 else '%rbp'}
            for lane,(start,i) in enumerate(zip(starts,comparisons))]


def main():
    out={'native':{},'limitations':'Bounded byte-addressed synthetic instruction evaluation, not original '
         'native/runtime/callee execution. Supplied statuses/fractions/payloads omit collision generation. '
         'Only finite comparisons and selected operations are modeled. Full entity-loop scheduling, '
         'geometry, malformed/exceptional inputs and other-platform 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 ['SV_ClipMoveToEntities','SV_Trace']};lanes=lane_sites(graphs['SV_ClipMoveToEntities'])
        clip,entity,stack=0x10000,0x20000,0x40000;current=clip+0x3c;candidate=stack+0x1030
        merges=[]
        for lane,status0,status4,fraction,prior4 in itertools.product(lanes,[0,1],[0,1],[.25,.5,.75],[0,1,4]):
            old=[0,prior4,bits(.5)]+[0x100+i for i in range(3,13)]+[42]
            new=[status0,status4,bits(fraction)]+[0x200+i for i in range(3,13)]+[999]
            m={current-1:0xab,current+56:0xcd};put_word(m,entity,7)
            for i,(v,w) in enumerate(zip(old,new)):put_word(m,current+4*i,v);put_word(m,candidate+4*i,w)
            t=evaluate(graphs['SV_ClipMoveToEntities'],lane['start'],[lane['stop']],m,
                       {'%rbx':clip,'%rsp':stack,lane['candidate_register']:entity})
            merges.append({'lane':lane['lane'],'candidate_status0':status0,'candidate_status4':status4,
                'candidate_fraction':fraction,'prior_status4':prior4,'old_words':old,'input_candidate_words':new,
                'final_words':[get_word(t['memory'],current+4*i) for i in range(14)],
                'canaries':[t['memory'][current-1],t['memory'][current+56]],'trace':t['trace'],'writes':t['writes']})
        g=graphs['SV_Trace'];rows=g['instructions']
        call=next(i for i,x in enumerate(rows) if '<CM_BoxTrace>' in x['operands'])
        start=rows[call+2]['address']  # Start after file-backed 1.0 load; supply xmm2 explicitly.
        branch=next(i for i in range(call,len(rows)) if rows[i]['opcode']=='je')
        stop_yes=int(rows[branch]['operands'].split()[0],16);stop_no=rows[branch+1]['address']
        world=[]
        for fraction in [0.0,-0.0,.25,1.0,1.5,-1.0]:
            m={};put_word(m,stack+0x64,bits(fraction))
            t=evaluate(g,start,[stop_yes,stop_no],m,{'%rsp':stack,'%xmm2':bits(1), '%rax':0})
            world.append({'fraction':fraction,'identifier_word':get_word(t['memory'],stack+0x90),
                          'skips_entity_clip_path':t['stop']==stop_yes,'trace':t['trace']})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],
            'graphs':graphs,'lane_sites':lanes,'composition_traces':merges,'world_decision_traces':world}
    write_json(Path('docs/research/collision-composition.json'),out)
    print('Recovered 216 complete three-lane merge/copy slices and 12 world-result decisions')


if __name__=='__main__':main()