#!/usr/bin/env python3
"""Independent existing momentum switch: native effects, gates, saves and menus.
Declared CPU/input fixtures and image controls are not human or win-rate trials.
"""
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import argparse,hashlib,json,os,struct,subprocess
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('--seed',type=Path,required=True);ap.add_argument('--state',type=Path,required=True);ap.add_argument('--section',choices=['all','storage','menu','events','gates','entries','game','negative'],default='all');a=ap.parse_args();h=Path(__file__).resolve().parent;o=a.output.resolve();o.mkdir(parents=True,exist_ok=True);env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};sha=lambda p:hashlib.sha256(p.read_bytes()).hexdigest();assert sha(a.baseline)=='c32a79e34127757b52cd488c124c9247bf7b875088bac89956a6b7ca2b132cf6'

def asm(p,body,origin=0x2d5000):
 p.mkdir(parents=True,exist_ok=True);(p/'driver.s').write_text(('' if '.globl _start' in body else ' .text\n .globl _start\n_start:\n')+body+'\n');subprocess.run([str(h.parent/'m68kasm.sh'),str(p/'driver.s'),f'{origin:x}',str(p/'code.bin')],stdout=subprocess.DEVNULL,check=True);return(p/'code.bin').read_bytes()
def core(p,rom,commands,settings=None):
 (p/'commands.txt').write_text('\n'.join(commands)+'\n')
 with(p/'log').open('w') as log:subprocess.run(['nice','-n','19','python3',str(h/'test_core.py'),str(rom.resolve()),str(p/'commands.txt'),str(p)],env=env|(settings or {}),stdout=log,stderr=log,check=True,timeout=120)
def run(p,body,size,patches=(),profile='current'):
 p.mkdir(parents=True,exist_ok=True);b=bytearray((a.baseline if profile=='parent' else a.rom).read_bytes());code=asm(p,' move.w #0x2700,%sr\n movea.l #0xfffff900,%sp\n'+body+'\n move.l #0x46444f54,0xffff3ff0\nend: bra.s end');assert len(code)<0x4000 and b[0x2d5000:0x2d5000+len(code)]==b'\xff'*len(code);b[0x2d5000:0x2d5000+len(code)]=code;struct.pack_into('>I',b,4,0x2d5000)
 for patch in patches:
  at,data=patch[:2]
  expected=patch[2] if len(patch)==3 else bytes.fromhex('08f80002c2ea') if at==0x12562 else b'\xff'*len(data)
  assert b[at:at+len(data)]==expected,(hex(at),'instrumentation reservation');b[at:at+len(data)]=data
 (p/'rom.bin').write_bytes(b);core(p,p/'rom.bin',['1500 dump ff3ff0 4 done',f'1500 dump ff4000 {size} result','1501 exit'],{'GPGX_SRM':str(p/'battery.srm')} if (p/'battery.srm').exists() else None)
 rows={t[1]:bytes.fromhex(' '.join(t[2:])) for line in(p/'commands.out').read_text().splitlines() if(t:=line.split())};assert rows['done']==b'FDOT',('native return',p);return rows['result']
def oracle(x,y):
 if not(-70<=x<=70 and -206<=y<=206):return x,y
 points=[(70,206),(-70,206)] if y>=166 else [(70,-206),(-70,-206)] if y<=-166 else [(0,0),(70,65),(-70,65),(70,-65),(-70,-65)]
 return min(points,key=lambda p:(p[0]-x)**2+(p[1]-y)**2)
def snapshot(at):return[f' movem.l %d0-%d7/%a0-%a6,0x{at:x}',f' move.w %sr,0x{at+60:x}',f' move.l %sp,0x{at+62:x}']
def abi(raw):return dict(regs=list(struct.unpack_from('>15I',raw)),sr=int.from_bytes(raw[60:62],'big'),sp=int.from_bytes(raw[62:66],'big'))
def registers(x,y):
 regs=[0x22330000+i for i in range(8)]+[0xffff1000+256*i for i in range(7)];regs[0]=0x12340000|(x&65535);regs[1]=0x56780000|(y&65535)
 return [f' move.l #0x{v:x},%d{i}' for i,v in enumerate(regs[:8])]+[f' movea.l #0x{v:x},%a{i}' for i,v in enumerate(regs[8:])],regs
CANARY=bytes((i*37+(i>>8)*11+19)&255 for i in range(0x4850))
def copy_native():return [' lea 0x230000,%a0',' lea 0xffffb000,%a1',' move.w #0x4850/4-1,%d0','copy: move.l (%a0)+,(%a1)+',' dbra %d0,copy']
def capture_ram(at):return[' lea 0xffffb000,%a0',f' lea 0x{at:x},%a1',' move.w #0x4850/4-1,%d0','capture: move.l (%a0)+,(%a1)+',' dbra %d0,capture']
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 storage_suite():
 specs=[dict(slot=slot,tag=t,inv=t^255) for slot in range(4) for t in range(256)]+[dict(slot=slot,tag=t,inv=t) for slot in range(4) for t in [0,0x7a,0x7b,255]]
 def trial(batch):
  ss=specs[batch*16:batch*16+16];body=[' move.b #1,0xa130f1']
  for j,z in enumerate(ss):
   rec=0x209aa5+4*z['slot'];body += [f' move.w #{int(z["slot"]>0)},0xfffff06c',f' move.w #{z["slot"]},0xfffff06a',f' move.b #{z["tag"]},0x{rec:x}',f' move.b #{z["inv"]},0x{rec+2:x}',' move.w #99,0xfffff84a'];init,regs=registers(12,-34);body+=init+[' move.w #0x271f,%sr',' jsr 0x226000']+snapshot(0xffff4000+j*68)+[f' move.w 0xfffff84a,0x{0xffff4000+j*68+66:x}']
  raw=run(o/f'storage-load-{batch}','\n'.join(body),68*len(ss));rows=[]
  for j,z in enumerate(ss):
   _,regs=registers(12,-34);ctx=abi(raw[j*68:]);value=struct.unpack_from('>H',raw,j*68+66)[0];want=z['tag']&1 if z['tag'] in [0x7a,0x7b] and z['tag']^z['inv']==255 else 0;assert ctx==dict(regs=regs,sr=0x271f,sp=0xfffff900) and value==want,('record validation/fallback/full ABI',z,value,ctx);rows.append(z|dict(value=value))
  return rows
 with ThreadPoolExecutor(max_workers=3) as pool:loads=sum(pool.map(trial,range((len(specs)+15)//16)),[])
 writes=[]
 for slot in range(4):
  for value in [0,1,99]:
   p=o/f'storage-put-{slot}-{value}';p.mkdir(exist_ok=True);b=bytearray((i*37+19)&255 for i in range(32768));b[0x4d22:0x4d52]=bytes.fromhex('619e609f619e609f718e708f718e708f738c728d738c728d758a748b758a748b76897689768976897986788779867887');b[0x4d52:0x4d5a]=bytes.fromhex('7a857a857a857a85');b=bytearray(pack(b)[1::2]);(p/'battery.srm').write_bytes(pack(b));body=[' move.b #1,0xa130f1',f' move.w #{int(slot>0)},0xfffff06c',f' move.w #{slot},0xfffff06a',f' move.w #{value},0xfffff84a',' jsr 0x226004',' move.b #0,0xa130f1',' move.w 0xfffff84a,0xffff4000'];run(p,'\n'.join(body),2);after=(p/'battery.srm').read_bytes()[1::2];allowed={0x4d52+2*slot,0x4d53+2*slot,0x5ffe,0x5fff};diff=[i for i,(x,y) in enumerate(zip(b,after)) if x!=y];assert set(diff)<=allowed;want=0x7b if value==1 else 0x7a;assert after[0x4d52+2*slot:0x4d54+2*slot]==bytes([want,want^255]);assert after[0x5fff]==sum(after[0x4000:0x5ffe])%256 and after[0x5ffe]==after[0x5fff]^255;writes.append(dict(slot=slot,value=value,changed=diff))
 print('PASS',len(loads),'storage/full ABI contexts and12complete64KB separated writes',flush=True);return dict(loads=loads,writes=writes)

from momentum_scenarios import Scenarios
scenario=Scenarios(a,h,o,asm,core,pack,run,registers)
suites={'storage':storage_suite,'menu':scenario.menu,'events':lambda:scenario.events(),'gates':lambda:scenario.gates(),'entries':lambda:scenario.entries(),'game':lambda:scenario.game(),'negative':lambda:scenario.negative()}
sections={name:fn() for name,fn in suites.items() if a.section in ['all',name]}
(o/'results.json').write_text(json.dumps(dict(rom_sha256=sha(a.rom),baseline_sha256=sha(a.baseline),sections=sections),indent=2)+'\n')
