#!/usr/bin/env python3
"""Native physical controller-star colors, original goal palette and real DMA."""
from pathlib import Path
import argparse,json,hashlib,os,struct,subprocess
from concurrent.futures import ThreadPoolExecutor
from pad_colors import STAR_ART,STAR_KEY,recolor
ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('rom',type=Path);ap.add_argument('output',type=Path);ap.add_argument('--baseline',type=Path,required=True);ap.add_argument('--r9-reference',type=Path,required=True);ap.add_argument('--seed',type=Path,required=True);a=ap.parse_args();o=a.output.resolve();o.mkdir(exist_ok=True,parents=True);h=Path(__file__).resolve().parent;new=a.rom.read_bytes();old=a.baseline.read_bytes();sha=lambda b:hashlib.sha256(b).hexdigest()
assert sha(old)=='501da8aaf8fc1fb4f5fdd22f4d8f8d893a8fe06de9ec28cb4ecaf0a6e18297eb'
assert sha(a.r9_reference.read_bytes())=='443c96a23923e96529120cd19c486db1483101896787c76a8d2082e5836e6186'
expected_palette=bytearray(old[0x59924:0x599a4]);expected_palette[0x36:0x38]=bytes.fromhex('00e0');assert new[0x59924:0x599a4]==expected_palette,'only the requested P3 green entry may differ; P1/goal and every other fixed entry preserved'
for lo,hi in [(0x7c144,0x7c244),(0x9e624,0x9e6a4),(0x2ff500,0x2ff800)]:assert new[lo:hi]==old[lo:hi],('original star art/native DMA/crit changed',hex(lo))
assert new[STAR_ART:STAR_ART+128]==recolor(old[0x7c144:0x7c1c4],9,11) and old[STAR_ART:STAR_ART+128]==b'\xff'*128
assert STAR_ART==0x4004 and ((STAR_KEY^0x8000)-0x8000)*32+0x5de84==STAR_ART,'signed native key must resolve to exact low-ROM padding'
for at in [0x2085da,0x2085e2]:assert int.from_bytes(new[at:at+2],'big')==STAR_KEY and old[at:at+2]==bytes.fromhex('0f1a')
for at in [0x205744,0x20574c]:assert int.from_bytes(new[at:at+2],'big')==int.from_bytes(old[at:at+2],'big')|0x2000
for at in [0x2085fc,0x208604]:assert new[at:at+2]==old[at:at+2],'native P4 yellow palette retained'
table=0x5de7a+int.from_bytes(old[0x5de7e:0x5de82],'big');count=int.from_bytes(old[table+2:table+4],'big')//2-1
changed_frames=[]
for frame in range(1,count):
 st,en=struct.unpack_from('>HH',old,table+2*frame);assert new[table+2*frame:table+2*frame+4]==old[table+2*frame:table+2*frame+4]
 for at in range(table+st,table+en,8):
  before=old[at:at+8];after=new[at:at+8];assert before[:4]==after[:4] and before[7]==after[7]
  size=before[7];n=((size>>2&3)+1)*((size&3)+1);ko=int.from_bytes(before[4:6],'big',signed=True);kn=int.from_bytes(after[4:6],'big',signed=True);pixels_o=old[0x5de84+32*ko:0x5de84+32*(ko+n)];pixels_n=new[0x5de84+32*kn:0x5de84+32*(kn+n)]
  assert [(v>>4!=0,v&15!=0) for v in pixels_o]==[(v>>4!=0,v&15!=0) for v in pixels_n],('shape changed',frame,at)
  if pixels_o!=pixels_n or before!=after:
   assert frame in [0x187,0x34c],('unrelated frame changed',hex(frame))
   if frame==0x34c:assert pixels_n==recolor(pixels_o,10,11) and before[6:]==after[6:],'P3 identical star shape, native palette1'
   changed_frames.append(frame)
assert set(changed_frames)=={0x187,0x34c}
# Native shared DMA entry, including puck/crit, player bodies and marker glyphs.
cases=[dict(actor=x,source=v,red=r,xflag=k) for x in [0xffffb04a,0xffffb74a,0xffffbdb4,0xffffbdd0,0xffffbdec,0xffffbe08,0xffffbe24,0xffffbe40] for v in [0x4004,0x4044,0x4083,0x7c1c4,0x7c244,0x7c324,0x7c404,0x60000] for r in [0,1] for k in [0,1]]
def run(spec):
 batch,items=spec;p=o/str(batch);p.mkdir(exist_ok=True);body=[' .text',' .globl _start','_start:',' move.w #0x2700,%sr',' movea.l #0xFFF600,%sp']
 for j,z in enumerate(items):
  for i in range(8):body.append(f' move.l #0x{0x12340000+i:X},%d{i}')
  for i in range(7):body.append(f' movea.l #0x{z["actor"] if i==3 else 0xffff1800+256*i:X},%a{i}')
  body += [f' move.l #0x{(z["source"]-0x5de84)&0xffffffff:X},%d2',f' move.w #{z["red"]},0xFFF062',f' move.w #{16 if z["xflag"] else 0},%ccr',' jsr 0x2FF510',f' move.w %sr,0x{0xff4000+j*66:X}',f' movem.l %d0-%d7/%a0-%a6,0x{0xff4002+j*66:X}',f' move.l %sp,0x{0xff403e+j*66:X}']
 body+=[' move.l #0x434F4C52,0xFF3FF0','end: bra.s end'];(p/'driver.s').write_text('\n'.join(body)+'\n');subprocess.run([str(h.parent/'m68kasm.sh'),str(p/'driver.s'),'2D5000',str(p/'code.bin')],check=True,stdout=subprocess.DEVNULL);code=(p/'code.bin').read_bytes();assert len(code)<0x3000
 raw=[]
 for name,b in [('current',new),('parent',old)]:
  b=bytearray(b);assert b[0x2d5000:0x2d5000+len(code)]==b'\xff'*len(code);b[0x2d5000:0x2d5000+len(code)]=code;struct.pack_into('>I',b,4,0x2d5000);(p/(name+'.bin')).write_bytes(b);(p/'commands.txt').write_text(f'500 dump ff3ff0 4 done\n500 dump ff4000 {len(items)*66} regs\n501 exit\n');env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env['GPGX_DEV']='513'
  with(p/(name+'.log')).open('w') as f:subprocess.run(['python3',str(h/'test_core.py'),str(p/(name+'.bin')),str(p/'commands.txt'),str(p)],env=env,stdout=f,stderr=f,check=True,timeout=60)
  rr={l.split()[1]:bytes.fromhex(' '.join(l.split()[2:])) for l in (p/'commands.out').read_text().splitlines()};assert rr['done']==b'COLR';raw.append(rr['regs']);(p/(name+'-regs.bin')).write_bytes(raw[-1])
 for j,z in enumerate(items):
  n=bytearray(raw[0][j*66:(j+1)*66]);b=bytearray(raw[1][j*66:(j+1)*66]);assert n==b,(z,'native DMA register/CCR/stack changed')
 return len(items)
with ThreadPoolExecutor(max_workers=3) as pool:counts=list(pool.map(run,[(i,cases[i*24:(i+1)*24]) for i in range((len(cases)+23)//24)]))


# Explicit geometry fixture after an actual cold boot. All palettes, markers,
# DMA/cache decisions and screen pixels are produced by the native game.
def pack(b):
 for off in range(0,len(b),8192):b[off+0x1fff]=sum(b[off:off+0x1ffe])&255;b[off+0x1ffe]=b[off+0x1fff]^255
 raw=bytearray(2*len(b));raw[1::2]=b;return raw
def battery():
 base=bytearray(a.seed.read_bytes()[1:16384:2]);base[0xe20:0x1ffe]=bytes(0x11de);banks=[]
 for slot,length in [(1,84),(2,31),(3,7)]:
  b=bytearray(base);struct.pack_into('>4s6H',b,0x1000,b'SESN',0,0,0,length,slot,3);struct.pack_into('>4s4H',b,0x1ea0,b'S3SV',1,slot,7,1);b[0x1ef0:0x1ef6]=bytes.fromhex('504304fb00ff');b[0x1ff6:0x1ffe]=bytes.fromhex('5500000202005a00');b[0x1eee:0x1ef0]=bytes.fromhex('50af');banks.append(b)
 b=bytearray(banks[0])+b''.join(banks);b[0x2e18:0x2e20]=bytes.fromhex('525001fe435401fe');b[0x4d20:0x4d22]=bytes.fromhex('50af');return pack(b)
# Preserve the native bank-swap instructions and observe which physical pad
# is temporarily aliased at C320/C322. The release ROM has no such probe.
bankdir=o/'bank-observer';bankdir.mkdir(exist_ok=True);(bankdir/'banks.s').write_text((h/'testcases/manual_keeper_banks.s').read_text());subprocess.run([str(h.parent/'m68kasm.sh'),str(bankdir/'banks.s'),'2EF400',str(bankdir/'banks.bin')],stdout=subprocess.DEVNULL,check=True);bankcode=(bankdir/'banks.bin').read_bytes();observed={}
for profile,path in [('current',a.rom),('parent',a.baseline),('r9',a.r9_reference)]:
 b=bytearray(path.read_bytes());assert b[0x2ef400:0x2ef400+len(bankcode)]==b'\xff'*len(bankcode);b[0x2ef400:0x2ef400+len(bankcode)]=bankcode
 for at,target,expected in [(0xad66,0x2ef400,'3f38c32831f8c32cc328'),(0xadce,0x2ef440,'31f8c320c32431dfc320'),(0xae00,0x2ef480,'3f38c32a31f8c32ec32a'),(0xae68,0x2ef4c0,'31f8c322c32631dfc322')]:
  assert b[at:at+10]==bytes.fromhex(expected);b[at:at+10]=bytes.fromhex('4ef9')+struct.pack('>I',target)+bytes.fromhex('4e714e71')
 observed[profile]=bankdir/(profile+'.bin');observed[profile].write_bytes(b)

def render(spec):
 profile,team,end,owner=spec;p=o/f'render-{profile}-{team}-{end}-{owner}';p.mkdir(exist_ok=True);z=dict(team=team,end=end,owner=owner);(p/'spec.json').write_text(json.dumps(z));(p/'battery.srm').write_bytes(battery());(p/'commands.txt').write_text('1500 press 1 Start 5\n1900 press 1 Start 5\n3550 snap\n3551 exit\n');rom=observed[profile];env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env.update(GPGX_DEV='1281',GPGX_SRM=str(p/'battery.srm'),GPGX_HOOK=str(h/'testcases/marker_colors_hook.py'))
 with(p/'log').open('w') as f:subprocess.run(['python3',str(h/'test_core.py'),str(rom),str(p/'commands.txt'),str(p)],env=env,stdout=f,stderr=f,check=True,timeout=60)
 d=json.loads((p/'render.json').read_text());v=(p/'vram.bin').read_bytes();raw=(p/'raw.state').read_bytes();assert raw.startswith(b'GENPLUS-GX 1.7.6');cram=d['cram'];assert d['controls']==[(team-1)*6+i for i in range(4)] and d['sides']==[team]*4,('physical input map',profile,z,d['controls'],d['sides'])
 # Calibration requires the actual native SAT link chain atDC00 and resident
 # native star pixels; no presumed object/palette mapping alone can pass.
 sat=[];i=0
 for _ in range(80):
  row=struct.unpack_from('<4H',v,0xdc00+i*8);sat.append(row);i=row[1]&127
  if i==0:break
 assert i==0 and len(sat)>=8,('native committed SAT chain',profile,z)
 stars=[]
 for pad,frame in enumerate([0x186,0x187,0x34c,0x34d]):
  obj=next(x for x in d['objects'] if int(x['frame'],16)&0x7ff==frame);assert int(obj['flags'],16)&0x8000==0,('physical star hidden',profile,z,pad,obj);tile=(int(obj['tile'],16)&0x7ff)+3;art=v[32*tile:32*(tile+4)];art=bytes(sum(([art[i+1],art[i]] for i in range(0,len(art),2)),[]));indices={x for b in art for x in [b>>4,b&15] if x};assert len(indices)==1,('star pattern corrupted/uncommitted',profile,z,pad,indices);index=next(iter(indices));rows=[q for q in sat if q[2]&0x7ff==tile]
  assert len(rows)==2 and all(q[1]>>8==5 for q in rows),('native star not actually in committed SAT',profile,z,pad,tile,rows)
  palettes={q[2]>>13&3 for q in rows};assert len(palettes)==1;palette=next(iter(palettes));rgbcode=cram[16*palette+index];rgb=[rgbcode&7,rgbcode>>3&7,rgbcode>>6&7];stars.append(dict(pad=pad+1,frame=hex(frame),object=obj['address'],palette=palette,index=index,RGB3=rgb,visible_native_sprites=len(rows),resident_pattern_sha256=sha(art)))
  if profile=='current':
   assert [palette,index]==[[1,9],[1,10],[1,11],[0,12]][pad],('wrong physical star color source',z,pad,palette,index)
   assert rgb==[[0,0,3],[7,0,0],[0,7,0],[6,5,1]][pad],('physical P1blue/P2red/P3green/P4yellow regression',z,pad,rgb)
   expected=recolor(old[0x7c144:0x7c1c4],9,index);assert art==expected,('native signed-key DMA did not deliver complete exact star geometry',z,pad)
 if profile in ['current','parent']:assert cram[16+9]==192 and cram[16+13]==3,'original blue P1 index and red goal index must remain restored'
 if profile=='r9':assert stars[0]['RGB3']==[7,0,0] and cram[16+13]==63,('reported r9 red-P1/yellow-goal palette regression not independently reproduced',z,stars,cram)
 return dict(profile=profile,spec=z,stars=stars,goal_palette_RGB3=[cram[29]&7,cram[29]>>3&7,cram[29]>>6&7],committed_SAT_sprites=len(sat),camera=d['camera'],screenshot=str((p/'0000.png').relative_to(o)))
with ThreadPoolExecutor(max_workers=3) as pool:renders=list(pool.map(render,[(profile,team,end,owner) for profile in ['current','parent','r9'] for team in [1,2] for end in [-1,1] for owner in range(4)]))
assert len(renders)==48
(o/'results.json').write_text(json.dumps(dict(rom_sha256=sha(new),baseline_sha256=sha(old),r9_sha256=sha(a.r9_reference.read_bytes()),native_DMA_contexts=sum(counts),all_native_frames_checked=count-1,changed_frames=sorted(set(changed_frames)),colors=['P1 blue','P2 red','P3 neon green','P4 native golden yellow'],all_fixed_palette_entries_except_requested_P3_green_exact=True,original_goal_palette_restored=True,no_runtime_render_hook=True,signed_sprite_art=dict(address=hex(STAR_ART),key=hex(STAR_KEY)),cold_boot_render_cases=renders,scope='Explicit post-boot geometry fixture; all native palettes/cache/DMA/committed SAT/resident star pixels execute. Both teams/ends/all four puck-owning physical pads. Native r8 controls and reported r9 red-P1/yellow-goal regression independently rendered. Not a human hardware playtest.'),indent=2)+'\n')
print('PASS',sum(counts),'paired native DMA ABI cases and',len(renders),'actual cold-boot physical-star/native-goal render cases',flush=True)