#!/usr/bin/env python3
"""Real native G/SOG hooks keep attacking keeper counters out of GA/SA."""
from pathlib import Path
import argparse,json,os,struct,subprocess
ap=argparse.ArgumentParser();ap.add_argument('rom',type=Path);ap.add_argument('output',type=Path);ap.add_argument('--allstar',action='store_true');a=ap.parse_args();h=Path(__file__).resolve().parent;o=a.output.resolve();o.mkdir(parents=True,exist_ok=True);rows=[]
cases=[dict(side=t,roster=i,kind=k) for t in (0,1) for i in (0,1,2,3,25) for k in ('goal','shot')]
cases += [dict(side=t,roster=1,kind=k,context=c) for t in (0,1) for k in ('goal','shot') for c in ('highlight','replay')]
for n,s in enumerate(cases):
 p=o/str(n);p.mkdir(exist_ok=True);side=s['side'];slot=s['roster'];base=0xffffc6ce+side*0x364;expected_keeper=slot<(3 if a.allstar or side==0 else 2) and not s.get('context');idx=3*side+slot if slot<3 else 0;body=[' .text',' .globl _start','_start:',' move.w #0x2700,%sr',' movea.l #0xfff900,%sp',' jsr 0x2df400',' move.w #10,0xffc330',' move.w #5,0xffc332',f' move.b #{16 if s.get("context")=="highlight" else 0},0xffc2ea',f' move.b #{128 if s.get("context")=="replay" else 0},0xffc2ec',f' move.w #{slot},0x{base+0x18:x}',' lea 0xffb04a,%a3',f' move.b #{slot},0x66(%a3)']
 if a.allstar:body+=[' move.w #1,0xfffff850',' move.w #1,0xfffff84c',' clr.w 0xfffff06c',' clr.w 0xffffc330',' move.w #1,0xffffc332']
 for reg in range(8):body.append(f' move.l #0x{0x11223300+reg:x},%d{reg}')
 for reg in range(7):body.append(f' movea.l #0x{base if reg==2 else 0xffb04a if reg==3 else 0xff6000 if reg==0 else 0xff1000+256*reg:x},%a{reg}')
 body += [' move.l #0x12345678,0xffd066',' move.w #0x2715,%sr',f' jsr 0x{0x14964 if s["kind"]=="goal" else 0x1301c:x}',' move.w %sr,0xff4800',' movem.l %d0-%d7/%a0-%a6,0xff4000',' move.l %sp,0xff403c',' move.w 0xff6000,0xff4810',' move.l 0xffd066,0xff4820',f' lea 0x{base:x},%a1',f' move.w #{slot if slot<3 else 0},%d1',f' move.w #{3+int(expected_keeper and s["kind"]=="goal")},%d2',f' move.w #{12+int(expected_keeper and s["kind"]=="shot")},%d3',' jsr 0x2df40c',' move.w %d2,0xff4830',' move.w %d3,0xff4832',' move.l #0x474b4154,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'),'2d9000',str(p/'code.bin')],check=True,stdout=subprocess.DEVNULL);data=bytearray(a.rom.read_bytes());code=(p/'code.bin').read_bytes();assert data[0x2d9000:0x2d9000+len(code)]==b'\xff'*len(code);data[0x2d9000:0x2d9000+len(code)]=code;struct.pack_into('>I',data,4,0x2d9000)
 # Observe exactly after displaced native instructions, before other accounting.
 for pc in (0x1496a,0x13022):data[pc:pc+2]=bytes.fromhex('4e75')
 (p/'driver.bin').write_bytes(data);(p/'commands.txt').write_text('100 dump ff3ff0 4 done\n100 dump ff4000 64 regs\n100 dump ff4800 52 observations\n100 dump fff800 14 counters\n101 exit\n');env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env['GPGX_AUTO_CREDITS']='0'
 with (p/'log.txt').open('w') as log:subprocess.run(['python3',str(h/'test_core.py'),str(p/'driver.bin'),str(p/'commands.txt'),str(p)],env=env,stdout=log,stderr=log,check=True,timeout=30)
 out={x.split()[1]:bytes.fromhex(' '.join(x.split()[2:])) for x in (p/'commands.out').read_text().splitlines()};assert out['done']==b'GKAT';regs=list(struct.unpack('>16I',out['regs']));want=[0x11223300+reg for reg in range(8)];want[0]=(0x11220000|slot);want += [base if reg==2 else 0xffb04a if reg==3 else 0xff6000+(s['kind']=='goal') if reg==0 else 0xff1000+256*reg for reg in range(7)]+[0xfff900];assert regs==want,(s,regs,want)
 counters=out['counters'];expected=bytearray(12)
 if expected_keeper:expected[2*idx+int(s['kind']=='shot')]=1
 assert counters[:12]==expected and counters[12:]==b'GK',(s,counters.hex(),expected.hex());obs=out['observations'];assert obs[32:36]==bytes.fromhex('12345678'),(s,'RNG changed');assert struct.unpack_from('>2H',obs,48)==(3,12),(s,'GA/SA subtraction',obs[48:].hex());ccr=int.from_bytes(obs[:2],'big')&31;assert ccr==(16|(4 if slot==0 else 0)),(s,'native MOVE CCR',ccr)
 if s['kind']=='goal':assert obs[16]==slot,(s,'native scorer log')
 rows.append(dict(case=s,keeper_attack_recorded=bool(expected_keeper),native_register_stack_CCR_RNG_preserved=True,defensive_GA_SA=[3,12]))
(o/'results.json').write_text(json.dumps(rows,indent=2)+'\n');print('PASS',len(rows),'native keeper attack / defensive accounting, skater exclusions, replay/highlight and complete ABI cases')