#!/usr/bin/env python3
from pathlib import Path
import struct, sys, re, json, hashlib, csv, os

ROOT=Path(__file__).resolve().parent.parent
OLD=Path(os.environ.get('PSNOVA_OLD_PATCH_CPK', ROOT/'NOVA_FileList_Vita_patch.cpk'))
COMPILED=Path(os.environ.get('PSNOVA_COMPILED_DIR', ROOT/'compiled'))
OUT=Path(os.environ.get('PSNOVA_RUNTIME_CPK_OUT', ROOT/'NOVA_FileList_Vita_patch_COMPLETE.cpk'))
ALIGN=8
TOC_OFFSET=0x800
RMD_RE=re.compile(r'^msg_(.+)\.rmd$',re.I)

sys.path.insert(0,str(Path(__file__).resolve().parent))
from canonical_source_lib import CPKIndex

TYPE_FMT={0x0:'>B',0x1:'>b',0x2:'>H',0x3:'>h',0x4:'>I',0x5:'>i',0x6:'>Q',0x7:'>q',0x8:'>f',0x9:'>d'}
TYPE_SIZE={t:struct.calcsize(f) for t,f in TYPE_FMT.items()}
TYPE_SIZE.update({0xA:4,0xB:8})

class StringPool:
    def __init__(self): self.buf=bytearray(); self.map={}
    def add(self,s):
        s='' if s is None else str(s)
        if s in self.map:return self.map[s]
        off=len(self.buf); self.map[s]=off; self.buf+=s.encode('utf-8')+b'\0'; return off

def build_utf(table_name, columns, rows):
    """columns: list of dict{name,type,storage,const}; storage 0x10/0x30/0x50."""
    sp=StringPool(); table_name_off=sp.add(table_name)
    for c in columns: sp.add(c['name'])
    # Add string constants/rows deterministically before encoding.
    for c in columns:
        if c['type']==0xA and c['storage']==0x30: sp.add(c.get('const',''))
    for row in rows:
        for c in columns:
            if c['storage']==0x50 and c['type']==0xA: sp.add(row.get(c['name'],'') or '')
    data=bytearray(); data_map={}
    def dataref(b):
        b=bytes(b or b'')
        if b in data_map:return data_map[b],len(b)
        off=len(data); data_map[b]=off; data.extend(b); return off,len(b)
    def enc_val(typ,v):
        if typ in TYPE_FMT:return struct.pack(TYPE_FMT[typ],0 if v is None else v)
        if typ==0xA:return struct.pack('>I',sp.add(v or ''))
        if typ==0xB:
            off,n=dataref(v or b''); return struct.pack('>II',off,n)
        raise ValueError(typ)
    coldef=bytearray()
    for c in columns:
        coldef.append(c['storage']|c['type']); coldef+=struct.pack('>I',sp.add(c['name']))
        if c['storage']==0x30: coldef+=enc_val(c['type'],c.get('const'))
    row_width=sum(TYPE_SIZE[c['type']] for c in columns if c['storage']==0x50)
    rows_blob=bytearray()
    for row in rows:
        start=len(rows_blob)
        for c in columns:
            if c['storage']==0x50: rows_blob+=enc_val(c['type'],row.get(c['name']))
        if len(rows_blob)-start!=row_width: raise AssertionError('row width mismatch')
    rows_off=0x18+len(coldef)
    strings_off=rows_off+len(rows_blob)
    strings_blob=bytes(sp.buf)
    data_off=strings_off+len(strings_blob)
    body=bytearray(struct.pack('>IIIIHHI',rows_off,strings_off,data_off,table_name_off,len(columns),row_width,len(rows)))
    body+=coldef+rows_blob+strings_blob+data
    return b'@UTF'+struct.pack('>I',len(body))+body

def chunk(magic, utf, padded=None):
    raw=magic+struct.pack('<III',0xFF,len(utf),0)+utf
    if padded is not None:
        if len(raw)>padded: raise ValueError('chunk exceeds padded size')
        raw+=b'\0'*(padded-len(raw))
    return raw

def align_up(n,a=ALIGN): return n if n%a==0 else n+(a-n%a)

def custom_sort(s): return ''.join('~' if c=='_' else c.lower() for c in s)

def sha256(path):
    h=hashlib.sha256()
    with open(path,'rb') as f:
        while True:
            b=f.read(8<<20)
            if not b:break
            h.update(b)
    return h.hexdigest()

old=CPKIndex(OLD)
old_by_name={e.filename:e for e in old.entries}
compiled_by_name={p.name:p for p in COMPILED.glob('msg_*.rmd')}
if len(compiled_by_name)!=24850: raise SystemExit(f'expected 24850 compiled, got {len(compiled_by_name)}')

# Preserve every old patch entry, overwrite any compiled RMD, add compiled RMDs absent from old patch.
names=set(old_by_name)|set(compiled_by_name)
names=sorted(names,key=custom_sort)
entries=[]
for name in names:
    if name in compiled_by_name:
        p=compiled_by_name[name]; sz=p.stat().st_size
        entries.append({'name':name,'kind':'compiled','path':p,'file_size':sz,'extract_size':sz})
    else:
        e=old_by_name[name]
        entries.append({'name':name,'kind':'oldraw','entry':e,'file_size':e.file_size,'extract_size':e.extract_size})

# TOC schema matches Nova public patch where possible.
toc_cols=[
 {'name':'DirName','type':0xA,'storage':0x30,'const':''},
 {'name':'FileName','type':0xA,'storage':0x50},
 {'name':'FileSize','type':0x4,'storage':0x50},
 {'name':'ExtractSize','type':0x4,'storage':0x50},
 {'name':'FileOffset','type':0x6,'storage':0x50},
 {'name':'ID','type':0x4,'storage':0x50},
 {'name':'UserString','type':0xA,'storage':0x30,'const':'<NULL>'},
]
rows=[{'FileName':e['name'],'FileSize':e['file_size'],'ExtractSize':e['extract_size'],'FileOffset':0,'ID':i} for i,e in enumerate(entries)]
toc0=build_utf('CpkTocInfo',toc_cols,rows)
toc_chunk_padded=align_up(16+len(toc0),ALIGN)
file_off=toc_chunk_padded
for row,e in zip(rows,entries):
    row['FileOffset']=file_off
    e['file_offset']=file_off
    file_off+=align_up(e['file_size'])
toc_utf=build_utf('CpkTocInfo',toc_cols,rows)
if align_up(16+len(toc_utf),ALIGN)!=toc_chunk_padded: raise AssertionError('TOC size changed after offsets')
content_offset=TOC_OFFSET+toc_chunk_padded
content_size=sum(align_up(e['file_size']) for e in entries)
etoc_offset=content_offset+content_size

# ETOC, N+1 rows, matching common CRI layout.
etoc_cols=[
 {'name':'UpdateDateTime','type':0x6,'storage':0x50},
 {'name':'LocalDir','type':0xA,'storage':0x30,'const':''},
]
etoc_rows=[{'UpdateDateTime':569143605520244480} for _ in entries]+[{'UpdateDateTime':0}]
etoc_utf=build_utf('CpkEtocInfo',etoc_cols,etoc_rows)
etoc_raw=chunk(b'ETOC',etoc_utf)
etoc_size=len(etoc_raw)

packed_sum=sum(e['file_size'] for e in entries)
extract_sum=sum(e['extract_size'] for e in entries)
enabled_packed=2*extract_sum
enabled_data=packed_sum+extract_sum
# Copy public patch schema, update dynamic values. Null fields remain null storage.
hdr_cols=[
 {'name':'UpdateDateTime','type':0x6,'storage':0x50},
 {'name':'FileSize','type':0x6,'storage':0x10},
 {'name':'ContentOffset','type':0x6,'storage':0x50},
 {'name':'ContentSize','type':0x6,'storage':0x50},
 {'name':'TocOffset','type':0x6,'storage':0x50},
 {'name':'TocSize','type':0x6,'storage':0x50},
 {'name':'TocCrc','type':0x4,'storage':0x10},
 {'name':'HtocOffset','type':0x6,'storage':0x10}, {'name':'HtocSize','type':0x6,'storage':0x10},
 {'name':'EtocOffset','type':0x6,'storage':0x50}, {'name':'EtocSize','type':0x6,'storage':0x50},
 {'name':'ItocOffset','type':0x6,'storage':0x10}, {'name':'ItocSize','type':0x6,'storage':0x10}, {'name':'ItocCrc','type':0x4,'storage':0x10},
 {'name':'GtocOffset','type':0x6,'storage':0x10}, {'name':'GtocSize','type':0x6,'storage':0x10}, {'name':'GtocCrc','type':0x4,'storage':0x10},
 {'name':'HgtocOffset','type':0x6,'storage':0x10}, {'name':'HgtocSize','type':0x6,'storage':0x10},
 {'name':'EnabledPackedSize','type':0x6,'storage':0x50}, {'name':'EnabledDataSize','type':0x6,'storage':0x50},
 {'name':'TotalDataSize','type':0x6,'storage':0x10}, {'name':'Tocs','type':0x4,'storage':0x10},
 {'name':'Files','type':0x4,'storage':0x50}, {'name':'Groups','type':0x4,'storage':0x50}, {'name':'Attrs','type':0x4,'storage':0x50},
 {'name':'TotalFiles','type':0x4,'storage':0x10}, {'name':'Directories','type':0x4,'storage':0x10}, {'name':'Updates','type':0x4,'storage':0x10},
 {'name':'Version','type':0x2,'storage':0x50}, {'name':'Revision','type':0x2,'storage':0x50}, {'name':'Align','type':0x2,'storage':0x50},
 {'name':'Sorted','type':0x2,'storage':0x50}, {'name':'EnableFileName','type':0x2,'storage':0x50}, {'name':'EID','type':0x2,'storage':0x10},
 {'name':'CpkMode','type':0x4,'storage':0x50}, {'name':'Tvers','type':0xA,'storage':0x50}, {'name':'Comment','type':0xA,'storage':0x10},
 {'name':'Codec','type':0x4,'storage':0x50}, {'name':'DpkItoc','type':0x4,'storage':0x50},
 {'name':'EnableTocCrc','type':0x2,'storage':0x50}, {'name':'EnableFileCrc','type':0x2,'storage':0x50}, {'name':'CrcMode','type':0x4,'storage':0x50},
 {'name':'CrcTable','type':0xB,'storage':0x50},
]
hrow={
 'UpdateDateTime':1,'ContentOffset':content_offset,'ContentSize':content_size,
 'TocOffset':TOC_OFFSET,'TocSize':16+len(toc_utf),'EtocOffset':etoc_offset,'EtocSize':etoc_size,
 'EnabledPackedSize':enabled_packed,'EnabledDataSize':enabled_data,'Files':len(entries),'Groups':0,'Attrs':0,
 'Version':7,'Revision':14,'Align':ALIGN,'Sorted':1,'EnableFileName':1,'CpkMode':1,
 'Tvers':old.header.get('Tvers') or 'CPKMC2.49.32, DLL3.24.00','Codec':0,'DpkItoc':0,
 'EnableTocCrc':0,'EnableFileCrc':0,'CrcMode':0,'CrcTable':b''
}
hdr_utf=build_utf('CpkHeader',hdr_cols,[hrow])
if 16+len(hdr_utf)>0x7FA: raise SystemExit('header too large')
hdr=bytearray(0x800); hdr[:4]=b'CPK '; struct.pack_into('<III',hdr,4,0xFF,len(hdr_utf),0); hdr[16:16+len(hdr_utf)]=hdr_utf; hdr[-6:]=b'(c)CRI'
toc_raw=chunk(b'TOC ',toc_utf,padded=toc_chunk_padded)

print('Building',OUT)
print('entries',len(entries),'old preserved names',len(old_by_name),'compiled replacements/additions',len(compiled_by_name),'new compiled',len(set(compiled_by_name)-set(old_by_name)))
print('toc',len(toc_raw),'content_offset',content_offset,'content_size',content_size,'etoc',etoc_offset,etoc_size)
with open(OUT,'wb') as out, open(OLD,'rb') as oldf:
    out.write(hdr); out.write(toc_raw)
    for i,e in enumerate(entries):
        if i%2000==0: print('write',i,'/',len(entries),e['name'])
        if e['kind']=='compiled':
            with open(e['path'],'rb') as f:
                while True:
                    b=f.read(4<<20)
                    if not b:break
                    out.write(b)
        else:
            oe=e['entry']; oldf.seek(oe.absolute_offset); remain=oe.file_size
            while remain:
                b=oldf.read(min(remain,4<<20))
                if not b: raise IOError('short old read '+e['name'])
                out.write(b); remain-=len(b)
        pad=align_up(e['file_size'])-e['file_size']
        if pad:out.write(b'\0'*pad)
    if out.tell()!=etoc_offset: raise AssertionError((out.tell(),etoc_offset))
    out.write(etoc_raw)

manifest={
 'output':str(OUT),'files':len(entries),'old_patch_files_preserved_or_replaced':len(old_by_name),
 'compiled_rmds':len(compiled_by_name),'compiled_rmds_existing_patch':len(set(compiled_by_name)&set(old_by_name)),
 'compiled_rmds_added_from_app_only':len(set(compiled_by_name)-set(old_by_name)),
 'content_offset':content_offset,'content_size':content_size,'toc_offset':TOC_OFFSET,'toc_size':16+len(toc_utf),
 'etoc_offset':etoc_offset,'etoc_size':etoc_size,'align':ALIGN,'packed_payload_sum':packed_sum,'extract_payload_sum':extract_sum,
 'size':OUT.stat().st_size,'sha256':sha256(OUT)
}
json.dump(manifest,open(ROOT/'runtime_cpk_build.json','w'),indent=2)
print(json.dumps(manifest,indent=2))
