"""Resolve actual original effect shader aliases by the native loader's keys.

4FA720 indexes ahsd3f headers by five uints at+10, then resolves the second
frozen effect object's20-byte key table(+58/+60) into the first object's shader
slots. Similar names, nearby DXBC and matching signatures are not substitutes.
"""
from pathlib import Path
import ctypes,hashlib,json,re,struct

folder=Path('artifacts/galaxy-query-research/power-outer-sources')
source=folder/'effects/Effects_Win64_SM50.arc.raw'
data=source.read_bytes()
if hashlib.sha256(data).hexdigest()!='649e3107c868fe31ba23c9e1d19be2b99777b908431987f04f08cd39f0f798ca':
    raise ValueError('Pinned shader archive changed')
out=folder/'authored-models/effects';out.mkdir(exist_ok=True)
shader_map={}
for match in re.finditer(re.escape(b'\x05\x02ahsd3f'),data):
    offset=match.start();size=struct.unpack_from('<I',data,offset+8)[0]
    if not 48<=size<=1024*1024 or offset+size>len(data):raise ValueError('Original shader record exceeds bounds')
    record=data[offset:offset+size]
    key=struct.unpack_from('<5I',record,16)
    names=list(dict.fromkeys(x[:-1].decode('ascii') for x in re.findall(rb'(?:vs|ps|gs|cs)_[A-Za-z0-9_]+_Win64_SM50\0',record)))
    if len(names)!=1 or key in shader_map:raise ValueError('Original shader key/name is ambiguous')
    positions=[]
    for item in re.finditer(b'DXBC',record):
        p=item.start()
        if p+32>len(record):continue
        n,chunks=struct.unpack_from('<2I',record,p+24)
        if 32<=n<=len(record)-p and 1<=chunks<=16:positions.append((p,n))
    if len(positions)!=1:raise ValueError('Original shader record container is ambiguous')
    p,n=positions[0];blob=record[p:p+n]
    shader_map[key]=dict(name=names[0],archiveOffset=offset,archiveBytes=size,dxbcOffset=offset+p,
        dxbcBytes=n,stage=struct.unpack_from('<H',record,0x26)[0],sha256=hashlib.sha256(blob).hexdigest(),blob=blob)

targets={'Alpha_UnlitNoScale_StarMapFacingAlpha_NearFade_Forward','DistantMilkyWayDust_Forward',
    'DistantMilkyWay_Forward','GalaxyGlow_Forward'}
rows=[]
dll=ctypes.WinDLL('d3dcompiler_47.dll');disassemble=dll.D3DDisassemble
disassemble.argtypes=[ctypes.c_void_p,ctypes.c_size_t,ctypes.c_uint,ctypes.c_char_p,ctypes.POINTER(ctypes.c_void_p)]
disassemble.restype=ctypes.c_int32
for match in re.finditer(re.escape(b'\x08\x02xfxd3f'),data):
    effect=match.start();cursor=effect+8;objects=[]
    for _ in range(2):
        marker=data.index(b'\0\1voMrtP',cursor)
        size,relocations=struct.unpack_from('<2I',data,marker+8)
        objects.append((marker+16,data[marker+16:marker+16+size]))
        cursor=marker+16+size+relocations*4
    first=objects[0][1];second=objects[1][1]
    pointer=struct.unpack_from('<Q',first)[0]
    name=first[pointer:first.index(b'\0',pointer)].decode('ascii')
    if name not in targets:continue
    pointer,count=struct.unpack_from('<2Q',second,0x58)
    first_count=struct.unpack_from('<Q',first,0x50)[0]
    if count!=first_count or count!=2:raise ValueError('Original model shader slot count changed')
    shaders=[]
    for index in range(count):
        key_offset=pointer+index*20;key=struct.unpack_from('<5I',second,key_offset)
        if key not in shader_map:raise ValueError('Original model shader key not found')
        original=shader_map[key];shader={k:v for k,v in original.items() if k!='blob'}
        shader.update(slot=index,key=[hex(x) for x in key],keyArchiveOffset=objects[1][0]+key_offset)
        shaders.append(shader);blob=original['blob']
        (out/(shader['name']+'.dxbc')).write_bytes(blob)
        result=ctypes.c_void_p();buffer=ctypes.create_string_buffer(blob)
        if disassemble(buffer,len(blob),0,None,ctypes.byref(result))<0:raise ValueError('Original linked shader disassembly failed')
        table=ctypes.cast(result,ctypes.POINTER(ctypes.POINTER(ctypes.c_void_p)))[0]
        try:
            p=ctypes.WINFUNCTYPE(ctypes.c_void_p,ctypes.c_void_p)(table[3])(result)
            n=ctypes.WINFUNCTYPE(ctypes.c_size_t,ctypes.c_void_p)(table[4])(result)
            text=ctypes.string_at(p,n).decode().rstrip('\0')
        finally:ctypes.WINFUNCTYPE(ctypes.c_uint,ctypes.c_void_p)(table[2])(result)
        (out/(shader['name']+'.txt')).write_text(text,encoding='utf-8')
    rows.append(dict(effect=name,effectOffset=effect,firstObjectOffset=objects[0][0],secondObjectOffset=objects[1][0],shaders=shaders))
if {row['effect'] for row in rows}!=targets:raise ValueError('Original effect link selection incomplete')
native=folder/'background/effect-loader/solo-4fa720.c'
result=dict(source=str(source),sourceSha256=hashlib.sha256(data).hexdigest(),shaderRecordCount=len(shader_map),effects=rows,
    nativeLoader=dict(imageSha256='6891f70a021fa9f3a0e00458c284f48d1226e68ae7fc02ad48081aed47857e93',
        rva='0x4FA720',source=str(native),sha256=hashlib.sha256(native.read_bytes()).hexdigest(),
        shaderKeyFields='+10,+14,+18,+1C,+20',effectKeyTable='secondFrozen+58/+60,20-byte stride',
        destinationSlots='firstFrozen+48,8-byte stride'),gameProcessAccess=False,gameCalls=0)
(out/'original-effect-shader-links.json').write_text(json.dumps(result,indent=2)+'\n',encoding='utf-8')
print(json.dumps({row['effect']:[s['name'] for s in row['shaders']] for row in rows}))