#!/usr/bin/env python3
from pathlib import Path
import argparse,csv,hashlib,struct,sys
from canonical_source_lib import CPKIndex,read_utf_table,cri_utf_decrypt

ALIGN=8; TOC_OFFSET=0x800
EXPECTED_APP_SHA='ad85c330c3258b0f18f57dcf77736f9fbea89c40b66dfef3fbb5c4daeff77e7a'
EXPECTED_APP_SIZE=2834095024
EXPECTED_OVERLAY_SHA='887a32c15927c227e989366f28f1624c95d7dbef9d8b781c4a6266040a160d17'
EXPECTED_OUTPUT_SHA='ebf2cea4c8158f32557d52c58b111101c1818157d862cd7b1c570b096730e4ba'
EXPECTED_OUTPUT_SIZE=3252228483
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})

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_file(p,bs=16<<20):
 h=hashlib.sha256()
 with open(p,'rb') as f:
  while True:
   b=f.read(bs)
   if not b: break
   h.update(b)
 return h.hexdigest()
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]
  o=len(self.buf);self.map[s]=o;self.buf+=s.encode()+b'\0';return o

def build_utf(name,cols,rows):
 sp=StringPool();tn=sp.add(name)
 for c in cols:sp.add(c['name'])
 for c in cols:
  if c['type']==0xA and c['storage']==0x30:sp.add(c.get('const',''))
 for r in rows:
  for c in cols:
   if c['storage']==0x50 and c['type']==0xA:sp.add(r.get(c['name'],'') or '')
 data=bytearray();dm={}
 def dataref(b):
  b=bytes(b or b'')
  if b in dm:return dm[b],len(b)
  o=len(data);dm[b]=o;data.extend(b);return o,len(b)
 def ev(t,v):
  if t in TYPE_FMT:return struct.pack(TYPE_FMT[t],0 if v is None else v)
  if t==0xA:return struct.pack('>I',sp.add(v or ''))
  if t==0xB:
   o,n=dataref(v or b'');return struct.pack('>II',o,n)
  raise ValueError(t)
 cd=bytearray()
 for c in cols:
  cd.append(c['storage']|c['type']);cd+=struct.pack('>I',sp.add(c['name']))
  if c['storage']==0x30:cd+=ev(c['type'],c.get('const'))
 rw=sum(TYPE_SIZE[c['type']] for c in cols if c['storage']==0x50);rb=bytearray()
 for r in rows:
  for c in cols:
   if c['storage']==0x50:rb+=ev(c['type'],r.get(c['name']))
 ro=0x18+len(cd);so=ro+len(rb);sb=bytes(sp.buf);do=so+len(sb)
 body=bytearray(struct.pack('>IIIIHHI',ro,so,do,tn,len(cols),rw,len(rows)));body+=cd+rb+sb+data
 return b'@UTF'+struct.pack('>I',len(body))+body

def chunk(magic,utf):
 enc=cri_utf_decrypt(utf)
 return magic+struct.pack('<III',0,len(enc),0)+enc

def etoc_rows(cpk):
 with cpk.path.open('rb') as f:r,_=read_utf_table(f,int(cpk.header['EtocOffset'])+16)
 return r

def main():
 ap=argparse.ArgumentParser(description='Build the verified Phantasy Star Nova Vita3K full English CPK from the clean APP CPK + Pass 11 overlay CPK.')
 ap.add_argument('--app',required=True,help='Original NOVA_FileList_Vita.cpk from the working Japanese Vita3K install')
 ap.add_argument('--overlay',required=True,help='NOVA_FileList_Vita_patch.cpk from PSNova_English_Runtime_Test_v0.1')
 ap.add_argument('--manifest',default=str(Path(__file__).with_name('compiled_manifest.csv')))
 ap.add_argument('--output',default='NOVA_FileList_Vita_FULL_ENGLISH.cpk')
 a=ap.parse_args();APP=Path(a.app);OV=Path(a.overlay);OUT=Path(a.output);MF=Path(a.manifest)
 print('Hashing inputs...')
 ash=sha256_file(APP);osh=sha256_file(OV)
 if APP.stat().st_size!=EXPECTED_APP_SIZE or ash!=EXPECTED_APP_SHA:raise SystemExit(f'Wrong APP CPK. Expected size {EXPECTED_APP_SIZE} SHA {EXPECTED_APP_SHA}; got {APP.stat().st_size} {ash}')
 if osh!=EXPECTED_OVERLAY_SHA:raise SystemExit(f'Wrong overlay CPK. Expected SHA {EXPECTED_OVERLAY_SHA}; got {osh}')
 app=CPKIndex(APP);ov=CPKIndex(OV)
 mf={}
 with MF.open(encoding='utf-8-sig',newline='') as f:
  for r in csv.DictReader(f):mf[str(r['ID'])]=r
 enabled=set(mf)
 if len(enabled)!=24850:raise SystemExit(f'Manifest must contain 24850 RMDs, got {len(enabled)}')
 if enabled-set(ov.rmd_by_id):raise SystemExit('Overlay is missing compiled RMDs')
 direct=enabled & set(app.rmd_by_id);patch_only=enabled-set(app.rmd_by_id)
 if len(direct)!=20930 or len(patch_only)!=3920:raise SystemExit(f'Unexpected mapping: direct={len(direct)} patch-only={len(patch_only)}')
 # Verify source RMD hashes before building.
 with OV.open('rb') as f:
  for i,rid in enumerate(sorted(enabled)):
   e=ov.rmd_by_id[rid];f.seek(e.absolute_offset);b=f.read(e.file_size)
   if e.file_size!=e.extract_size or hashlib.sha256(b).hexdigest()!=mf[rid]['sha256']:raise SystemExit('Overlay RMD mismatch: '+rid)
   if i and i%5000==0:print(' verified',i,'/24850 RMDs')
 app_et=etoc_rows(app);ov_et=etoc_rows(ov)
 additions=sorted(patch_only,key=lambda r:custom_sort('msg_'+r+'.rmd'));add_id={r:len(app.entries)+i for i,r in enumerate(additions)}
 append_abs=int(app.header['EtocOffset']);offs={}
 for rid in sorted(enabled,key=lambda r:custom_sort('msg_'+r+'.rmd')):
  e=ov.rmd_by_id[rid];offs[rid]=append_abs-TOC_OFFSET;append_abs+=align_up(e.file_size)
 new_etoc_offset=append_abs
 final=[]
 for e in app.entries:
  rid=e.filename[4:-4] if e.filename.lower().startswith('msg_') and e.filename.lower().endswith('.rmd') else None
  if rid in enabled:
   s=ov.rmd_by_id[rid];final.append({'n':e.filename,'fs':s.file_size,'xs':s.extract_size,'fo':offs[rid],'id':int(e.cpk_id),'rid':rid})
  else:final.append({'n':e.filename,'fs':e.file_size,'xs':e.extract_size,'fo':e.file_offset,'id':int(e.cpk_id),'rid':None})
 for rid in additions:
  e=ov.rmd_by_id[rid];final.append({'n':e.filename,'fs':e.file_size,'xs':e.extract_size,'fo':offs[rid],'id':add_id[rid],'rid':rid})
 final.sort(key=lambda e:custom_sort(e['n']))
 tc=[{'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>'}]
 tr=[{'FileName':e['n'],'FileSize':e['fs'],'ExtractSize':e['xs'],'FileOffset':e['fo'],'ID':e['id']} for e in final];toc=chunk(b'TOC ',build_utf('CpkTocInfo',tc,tr))
 co=int(app.header['ContentOffset'])
 if len(toc)>co-TOC_OFFSET:raise SystemExit('Expanded TOC does not fit before original content')
 er=[{'UpdateDateTime':int(app_et[i].get('UpdateDateTime') or 0)} for i in range(len(app.entries))]
 for rid in additions:
  e=ov.rmd_by_id[rid];er.append({'UpdateDateTime':int(ov_et[int(e.cpk_id)].get('UpdateDateTime') or 0)})
 er.append({'UpdateDateTime':0});ec=[{'name':'UpdateDateTime','type':0x6,'storage':0x50},{'name':'LocalDir','type':0xA,'storage':0x30,'const':''}];etoc=chunk(b'ETOC',build_utf('CpkEtocInfo',ec,er))
 ps=sum(e['fs'] for e in final);es=sum(e['xs'] for e in final)
 hc=[{'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':'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':'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':'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}]
 hr={'UpdateDateTime':int(app.header.get('UpdateDateTime') or 1),'ContentOffset':co,'ContentSize':new_etoc_offset-co,'TocOffset':TOC_OFFSET,'TocSize':len(toc),'EtocOffset':new_etoc_offset,'EtocSize':len(etoc),'EnabledPackedSize':2*es,'EnabledDataSize':ps+es,'Files':len(final),'Groups':0,'Attrs':0,'Version':7,'Revision':1,'Align':8,'Sorted':1,'CpkMode':1,'Tvers':app.header['Tvers'],'Codec':0,'DpkItoc':0}
 hu=build_utf('CpkHeader',hc,[hr]);he=cri_utf_decrypt(hu);hdr=bytearray(TOC_OFFSET);hdr[:4]=b'CPK ';struct.pack_into('<III',hdr,4,0,len(he),0);hdr[16:16+len(he)]=he;hdr[-6:]=b'(c)CRI'
 print('Building',OUT,'(about 3.25 GB)...')
 with OUT.open('wb') as out,APP.open('rb') as af,OV.open('rb') as of:
  out.write(hdr);out.write(toc);out.write(b'\0'*(co-out.tell()));af.seek(co);remain=int(app.header['EtocOffset'])-co
  while remain:
   b=af.read(min(remain,16<<20));out.write(b);remain-=len(b)
  for i,rid in enumerate(sorted(enabled,key=lambda r:custom_sort('msg_'+r+'.rmd'))):
   e=ov.rmd_by_id[rid];of.seek(e.absolute_offset);out.write(of.read(e.file_size));out.write(b'\0'*(align_up(e.file_size)-e.file_size))
   if i and i%5000==0:print(' appended',i,'/24850')
  out.write(etoc)
 oh=sha256_file(OUT);sz=OUT.stat().st_size
 print('Output size:',sz);print('SHA-256:',oh)
 if sz!=EXPECTED_OUTPUT_SIZE or oh!=EXPECTED_OUTPUT_SHA:raise SystemExit('BUILD FINISHED BUT HASH DOES NOT MATCH VERIFIED PASS 12 OUTPUT')
 print('PASS: output is byte-identical to verified Pass 12 build.')
if __name__=='__main__':main()
