"""Build local UI introspection overrides from this installation's assets.

Uses pyflashkit (MIT), zopfli and the locally checked out Conjunction packer
(own-use license). Original game bundles are read only. No game assets are
redistributed. Unknown GFx tags and CR2W payloads are preserved verbatim.
"""
from pathlib import Path
import hashlib
import json
import struct
import sys
import zlib

ROOT=Path(__file__).resolve().parent.parent

# Index bundle locations are relative to the configured game installation.
def indexed_bundle(value):
    path=Path(value)
    if path.is_absolute() or '..' in path.parts:raise ValueError('Nonportable UI index bundle path')
    return game_paths.GAME/path
from . import game_paths
sys.path[:0]=[str(game_paths.tool_path('repos/pyflashkit')),str(game_paths.tool_path('repos/Conjunction'))]
import zopfli.zlib
from flashkit.swf.parser import parse_swf
from flashkit.abc.parser import parse_abc,write_u30
from flashkit.abc.writer import serialize_abc
from flashkit.abc.types import MultinameInfo,NamespaceInfo,MethodInfo,MethodBodyInfo,TraitInfo
from conjunction import packer,metastore
from . import support

def patch(original):
    offset=original.index(b'CFX')
    capacity=struct.unpack_from('<I',original,offset-4)[0]
    stream=original[offset:offset+capacity]
    raw=b'FWS'+stream[3:8]+zlib.decompress(stream[8:])
    header,tags,_,_=parse_swf(raw)
    root=None
    for t in tags:
        if t.tag_type!=76:continue
        p=2
        for _ in range(struct.unpack_from('<H',t.payload)[0]):
            idx=struct.unpack_from('<H',t.payload,p)[0];p+=2
            end=t.payload.index(0,p);name=t.payload[p:end].decode();p=end+1
            if idx==0:root=name
    if not root:raise ValueError('No AVM2 document class')
    p=len(header)
    while p<len(raw):
        tag=struct.unpack_from('<H',raw,p)[0];size=tag&63;hs=2
        if size==63:size=struct.unpack_from('<I',raw,p+2)[0];hs=6
        if tag>>6==82:
            start=raw.index(b'\0',p+hs+4)+1
            before=raw[start:p+hs+size];a=parse_abc(before)
            def qualified(mn):
                m=a.multiname_pool[mn];ns=a.namespace_name(m.ns)
                return (ns+'.' if ns else '')+a.string(m.name)
            inst=next((i for i in a.instances if qualified(i.name)==root),None)
            if inst:
                assert hs==6 and serialize_abc(a)==before
                break
        p+=hs+size
    else:raise ValueError('Document class not defined in this movie')
    def string(v):
        if v not in a.string_pool:a.string_pool.append(v)
        return a.string_pool.index(v)
    def qname(v,package=''):
        ns=next((i for i,n in enumerate(a.namespace_pool) if n.kind==0x16 and a.string(n.name)==package),None)
        if ns is None:
            ns=len(a.namespace_pool);a.namespace_pool.append(NamespaceInfo(0x16,string(package)))
        a.multiname_pool.append(MultinameInfo(kind=7,ns=ns,name=string(v)))
        return len(a.multiname_pool)-1
    def op(code,*args):return bytes([code])+b''.join(write_u30(i) for i in args)
    scope=next(b.init_scope_depth for b in a.method_bodies if b.method==inst.iinit)
    def method(name,code,return_type='Object',params=(),getter=True,extra_locals=0):
        idx=len(a.methods)
        a.methods.append(MethodInfo(param_count=len(params),return_type=qname(return_type),param_types=[qname(t) for t in params],name=string(name),flags=0))
        a.method_bodies.append(MethodBodyInfo(method=idx,max_stack=16 if extra_locals else 8,local_count=1+len(params)+extra_locals,init_scope_depth=scope,max_scope_depth=scope+1,code=b'\xd0\x30'+code+b'\x48'))
        inst.traits.append(TraitInfo(name=qname(name),kind=2 if getter else 1,method_idx=idx))
    parent,index,other=qname('pdParent'),qname('pdIndex'),qname('pdOther')
    inst.traits.extend([TraitInfo(name=parent,kind=0,type_name=qname('Object')),TraitInfo(name=index,kind=0,type_name=qname('int')),TraitInfo(name=other,kind=0,type_name=qname('Object'))])
    getparent=b'\xd0'+op(0x66,parent)
    method('pdProbe',op(0x2c,string('PathDevkitUIv3')),'String')
    method('pdSame',getparent+b'\xd0'+op(0x66,other)+b'\xac','Boolean')
    method('pdAttached',getparent+op(0x66,qname('stage'))+b'\xd0'+op(0x66,qname('stage'))+b'\xac','Boolean')
    method('pdChild',getparent+b'\xd0'+op(0x66,index)+op(0x46,qname('getChildAt'),1))
    method('pdBounds',getparent+b'\xd0'+op(0x66,qname('stage'))+op(0x46,qname('getBounds'),1))
    classname=qname('getQualifiedClassName','flash.utils')
    method('pdClass',op(0x5d,classname)+getparent+op(0x46,classname,1),'String')
    event=qname('MouseEvent','flash.events')
    method('pdDispatch',getparent+op(0x5d,event)+b'\xd1\x26\x27'+op(0x4a,event,3)+op(0x46,qname('dispatchEvent'),1),'Boolean',('String',),False)
    if root == 'red.game.witcher3.hud.Hud':
        from .overlay_flash import LAYER, methods
        def uint(value):
            if value not in a.uint_pool:a.uint_pool.append(value)
            return a.uint_pool.index(value)
        inst.traits.append(TraitInfo(name=qname(LAYER),kind=0,type_name=qname('Object')))
        for declaration in methods(op,qname,string,uint):
            method(*declaration)
    after=serialize_abc(a)
    raw=raw[:p+2]+struct.pack('<I',size+len(after)-len(before))+raw[p+6:start]+after+raw[p+hs+size:]
    replacement=b'CFX'+raw[3:4]+struct.pack('<I',len(raw))+zopfli.zlib.compress(raw[8:])
    if len(replacement)>capacity:raise ValueError('Patched stream exceeds fixed CR2W allocation')
    out=bytearray(original);out[offset:offset+capacity]=replacement+bytes(capacity-len(replacement))
    table,count,_=struct.unpack_from('<III',out,88)
    for i in range(count):
        entry=table+24*i;length,start=struct.unpack_from('<II',out,entry+8)
        assert zlib.crc32(original[start:start+length])==struct.unpack_from('<I',original,entry+20)[0]
        struct.pack_into('<I',out,entry+20,zlib.crc32(out[start:start+length]))
    struct.pack_into('<I',out,96,zlib.crc32(out[table:table+24*count]))
    struct.pack_into('<I',out,32,zlib.crc32(out[:32]+struct.pack('<I',0xdeadbeef)+out[36:160]))
    return bytes(out),{'root':root,'capacity':capacity,'used':len(replacement)}

def build(names):
    folder=ROOT/'research/ui-build';cooked=folder/'cooked';packed=folder/'packed'
    cooked.mkdir(parents=True,exist_ok=True)
    index=json.loads((ROOT/'research/all-ui-index.json').read_text())
    report=[]
    seen=set()
    for row in index:
        if row['name'] in seen:continue
        seen.add(row['name'])
        with indexed_bundle(row['bundle']).open('rb') as f:
            if names and not any(n in row['name'] for n in names):continue
            f.seek(row['offset']);raw=f.read(row['compressed'])
            assert row['codec'] in (0,1)
            if row['codec']==1:raw=zlib.decompress(raw)
            assert len(raw)==row['uncompressed']
            try:result,details=patch(raw)
            except (ValueError,StopIteration) as e:
                report.append({'path':row['name'],'supported':False,'error':str(e)});continue
            path=cooked/row['name'].replace('\\','/');path.parent.mkdir(parents=True,exist_ok=True);path.write_bytes(result)
            report.append({'path':row['name'],'supported':True,'original_sha256':hashlib.sha256(raw).hexdigest(),'patched_sha256':hashlib.sha256(result).hexdigest(),**details})
            print(row['name'],details,flush=True)
    bundle=packer.pack(str(cooked),str(packed))
    (packed/'metadata.store').write_bytes(metastore.build([('blob0.bundle',bundle)]))
    (folder/'report.json').write_text(json.dumps(report,indent=2)+'\n')
    return report

def build_overlay():
    """Replace only the HUD movie, preserving every other hash-owned UI resource."""
    from conjunction import bundles
    from . import lifecycle
    import time
    lifecycle.verify()
    receipt=json.loads((support.STATE_ROOT/'deployment.json').read_text())
    owned=next(Path(e['path']) for e in receipt['assets'] if Path(e['path']).name=='blob0.bundle')
    folder=ROOT/'research/ui-build'
    cooked=folder/('overlay-source-'+str(time.time_ns()))
    packed=folder/'packed'
    cooked.mkdir(parents=True)
    hud='gameplay\\gui_new\\swf\\hud\\hud.redswf'
    preserved={}
    for entry in bundles.entries(str(owned)):
        name=entry[0]
        relative=Path(name.replace('\\','/'))
        if relative.is_absolute() or '..' in relative.parts:raise ValueError('Unsafe UI depot path')
        if name==hud:continue
        raw=bundles.read(str(owned),entry)
        path=cooked/relative;path.parent.mkdir(parents=True,exist_ok=True);path.write_bytes(raw)
        preserved[name]=hashlib.sha256(raw).hexdigest()
    row=next(r for r in json.loads((ROOT/'research/all-ui-index.json').read_text()) if r['name']==hud)
    with indexed_bundle(row['bundle']).open('rb') as f:
        f.seek(row['offset']);raw=f.read(row['compressed'])
    if row['codec']==1:raw=zlib.decompress(raw)
    elif row['codec']!=0:raise ValueError('Unexpected HUD codec')
    result,details=patch(raw)
    path=cooked/Path(hud.replace('\\','/'));path.parent.mkdir(parents=True,exist_ok=True);path.write_bytes(result)
    bundle=packer.pack(str(cooked),str(packed))
    expected=dict(preserved)
    expected[hud]=hashlib.sha256(result).hexdigest()
    actual={e[0]:hashlib.sha256(bundles.read(str(packed/'blob0.bundle'),e)).hexdigest()
            for e in bundles.entries(str(packed/'blob0.bundle'))}
    if actual!=expected:raise ValueError('Overlay bundle round-trip changed another UI resource')
    (packed/'metadata.store').write_bytes(metastore.build([('blob0.bundle',bundle)]))
    from .overlay_flash import VERSION
    report=dict(changed_depot_path=hud,preserved_resources=preserved,renderer_version=VERSION,
                original_hud_sha256=hashlib.sha256(raw).hexdigest(),
                patched_hud_sha256=hashlib.sha256(result).hexdigest(),**details)
    (folder/'overlay-report.json').write_text(json.dumps(report,indent=2)+'\n')
    print(json.dumps(dict(report,preserved_resources=len(preserved))))
    return report


if __name__=='__main__':
    if sys.argv[1:]==['--overlay']:build_overlay()
    else:build(sys.argv[1:])