#!/usr/bin/env python3
"""Real menu lock, cold-boot settings, legacy recovery and native first-game captures.
All batteries are generated/copied under the explicit private output directory.
"""
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import argparse,json,os,struct,subprocess
ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('rom',type=Path);ap.add_argument('output',type=Path);ap.add_argument('--seed',type=Path,required=True);a=ap.parse_args();here=Path(__file__).resolve().parent;out=a.output.resolve();out.mkdir(parents=True,exist_ok=True);reports=[]
F=[0xd050,0xd052,0xd056,0xd058,0xf004,0xf07a]
PRESETS=[[0,0,2,2,1,2],[1,1,1,0,0,1],[2,0,2,1,1,0]]

def put(p,b):
 for off in range(0,len(b),8192):b[off+0x1fff]=sum(b[off:off+0x1ffe])%256;b[off+0x1ffe]=b[off+0x1fff]^255
 raw=bytearray(len(b)*2);raw[1::2]=b;p.write_bytes(raw)
def settings(v,tag=0x54):
 check=0x5a
 for x in v[:5]:check^=x
 return bytes([tag,*v[:5],check])
def seed(path,kind='valid',gp=1):
 base=bytearray(a.seed.read_bytes()[1:16384:2]);base[0xe20:0x1ffe]=bytes(0x11de)
 banks=[]
 for slot in (1,2,3):
  b=bytearray(base);struct.pack_into('>4s6H',b,0x1000,b'SESN',gp,1,gp,[84,31,7][slot-1],0 if gp else slot,3)
  struct.pack_into('>4s4H',b,0x1ea0,b'S3SV',1,slot,7,1)
  v=PRESETS[slot-1]
  # Old slot snapshots contain four original choices, including D072=Auto.
  for addr,value in zip(F[:4],v[:4]):struct.pack_into('>H',b,0x1eac+addr-0xd046,value)
  struct.pack_into('>H',b,0x1eac+0xd072-0xd046,v[3])
  b[0x1ef0:0x1ef4]=b'PC\x01\xfe';b[0x1ef4:0x1ef6]=bytes([v[5],v[5]^255])
  if kind=='valid':b[0x1ff6:0x1ffd]=settings(v)
  elif kind=='old-auto':b[0x1ff6:0x1ffd]=settings([*v[:3],0 if v[3]==2 else v[3],*v[4:]],0x53)
  elif kind=='enum':b[0x1ff6:0x1ffd]=settings([99,*v[1:]])
  elif kind=='invalid':b[0x1ff6:0x1ffd]=settings(v);b[0x1ffc]^=1
  banks.append(b)
 put(path,bytearray(banks[0])+b''.join(banks))

hook=out/'hook.py';hook.write_text('''from pathlib import Path
import json,struct
p=Path(outdir);spec=json.loads((p/'spec.json').read_text());checks=[];injected=False
F=[0xd050,0xd052,0xd056,0xd058,0xf004,0xf07a]
def rw(a):return (wram[ram(a)]<<8)|wram[ram(a+1)]
def ww(a,v):wram[ram(a)]=v>>8;wram[ram(a+1)]=v&255
def emit(f):
 b=bytes(sram)[1:16384:2];row=dict(frame=f,slot=rw(0xf06a),view=rw(0xf06c),forced=rw(0xe3f0),values=[rw(x) for x in F],line_original=rw(0xd072),settings=b[0x1ff6:0x1ffd].hex(),difficulty=b[0x1ef4:0x1ef6].hex(),gp=int.from_bytes(b[0x1008:0x100a],'big'),clock=rw(0xc468))
 checks.append(row);(p/'observed.json').write_text(json.dumps(checks,indent=2)+'\\n')
def test_tick(f):
 global injected
 if spec.get('inject') and f==1000:
  for addr,v in zip(F,spec['inject']):ww(addr,v)
  ww(0xd072,spec['inject'][3])
 if spec.get('finish') and not injected and f>=3500 and rw(0xc468)>0 and not wram[ram(0xc2ea)]&0x11 and not wram[ram(0xc2ec)]&128 and not 0xa501<=rw(0xf738)<=0xa507:
  # Explicit leading-side goal-line fixture, as in the native shot-plan menu
  # suite. Native goals/GP/checksum/save must complete; never write headers.
  shooter=0 if rw(0xc6da)>=rw(0xca3e)else 6
  positive=bool(wram[ram(0xb04a+128*shooter+0x62)]&128)
  for actor in range(12):
   obj=0xb04a+128*actor;ww(obj,(140 if actor%2 else -140)&65535);ww(obj+0x14,0)
   for delta in [2,0x16,0x28,0x2a]:ww(obj+delta,0)
  ww(0xc466,3);ww(0xc468,600);ww(0xc46a,1);ww(0xbeda,shooter);ww(0xbee0,65535)
  for addr in [0xb74a,0xb74c,0xb760,0xb762,0xb764,0xb772,0xb776]:ww(addr,0)
  ww(0xb75e,260 if positive else (-260)&65535);ww(0xb774,0x1000 if positive else (-0x1000)&65535);ww(0xb7aa,65535);wram[ram(0xc2ea)]&=254;injected=True
 if spec.get('finish') and f>3500:
  b=bytes(sram)[1:16384:2];gp=int.from_bytes(b[0x1008:0x100a],'big')
  if gp==1 and rw(0xe3f0)==1 and rw(0xe3f2)==1 and b[0x1fff]==sum(b[:0x1ffe])%256:
   emit(f);assert checks[-1]['values']==spec['inject'],('completed-game menu changed settings',checks[-1])
   (p/'battery.srm').write_bytes(bytes(sram)[:65536]);print('PASS actual first game saved',f);sys.exit(0)
  if f%120==0 and (rw(0xc466)>=3 and wram[ram(0xc2ea)]&64 or wram[ram(0xc2ec)]&128):held[(0,IDS['Start'])]=5
 if str(f) in spec.get('checks',{}):
  expected=spec['checks'][str(f)];emit(f);r=checks[-1]
  assert r['slot']==expected['slot'] and r['view']==1,(spec['name'],r,expected)
  assert r['values']==expected['values'] and r['line_original']==expected['values'][3],(spec['name'],r,expected)
 if spec.get('capture') and f in (1200,3300):
  emit(f);r=checks[-1];wanted=spec['inject']
  assert bytes.fromhex(r['settings'])[:6]==bytes([0x55,*wanted[:5]]),('incorrect first-Start capture',r,wanted)
  assert bytes.fromhex(r['difficulty'])==bytes([wanted[5],wanted[5]^255]),('difficulty not captured',r,wanted)
  if f==3300:
   assert r['clock']>0 and rw(0xc466)<=3,('native match did not start',r)
   assert r['line_original']==wanted[3],('Auto menu value lost in-game',r,wanted)
''')

def run(name,battery,lines,checks=None,inject=None,capture=False,finish=False):
 p=out/name;p.mkdir(exist_ok=True);(p/'battery.srm').write_bytes(battery.read_bytes());(p/'spec.json').write_text(json.dumps(dict(name=name,checks=checks or {},inject=inject,capture=capture,finish=finish),indent=2)+'\n');(p/'commands.txt').write_text('\n'.join(lines)+'\n')
 env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env.update(GPGX_DEV='513',GPGX_SRM=str(p/'battery.srm'),GPGX_HOOK=str(hook))
 with (p/'log.txt').open('w') as log:subprocess.run(['nice','-n','19','python3',str(here/'test_core.py'),str(a.rom.resolve()),str(p/'commands.txt'),str(p)],env=env,check=True,stdout=log,stderr=log,timeout=35)
 observations=json.loads((p/'observed.json').read_text());b=(p/'battery.srm').read_bytes()[1::2];assert b[0x1fff]==sum(b[:0x1ffe])%256
 if finish:assert any(row['gp']==1 and row['forced']==1 for row in observations),'native first game did not actually commit GP1 within original13000-frame budget'
 return dict(test=name,checks=len(observations)),p/'battery.srm'

# Every settings row: both directions blocked in all three started season slots.
for kind in ['valid','legacy','invalid','enum','old-auto']:
 source=out/(kind+'.srm');seed(source,kind)
 values=[list(v) for v in PRESETS]
 if kind in ['legacy','invalid','enum']:
  for v in values:v[4]=0  # Fighting was not in the old snapshot; explicit legacy fallback.
 lines=[];checks={1000:dict(slot=1,values=values[0])};t=1100
 for slot in (1,2,3):
  for _ in range(4):lines.append(f'{t} press 1 Down 5');t+=90
  for row in range(6):
   for direction in ['Left','Right']:
    lines.append(f'{t} press 1 {direction} 5');checks[t+65]=dict(slot=slot,values=values[slot-1]);t+=90
   if row<5:lines.append(f'{t} press 1 Down 5');t+=90
  for _ in range(9):lines.append(f'{t} press 1 Up 5');t+=90
  lines.append(f'{t} press 1 Right 5');t+=120
 # Mini -> native modes -> Full, restore its choices despite native-mode defaults.
 for _ in range(5 if a.rom.read_bytes()[0x22c000:0x22c002]!=b'\xff\xff' else 4 if a.rom.read_bytes()[0x229000:0x229002]!=b'\xff\xff' else 3):lines.append(f'{t} press 1 Right 5');t+=120
 checks[t-30]=dict(slot=1,values=values[0]);lines.append(f'{t+50} exit')
 result,battery=run(kind+'-locked-circles',source,lines,checks);reports.append(result)
 checks={1000:dict(slot=1,values=values[0]),1260:dict(slot=2,values=values[1]),1460:dict(slot=3,values=values[2]),1660:dict(slot=2,values=values[1]),1860:dict(slot=1,values=values[0])}
 result,_=run(kind+'-cold-reboot',battery,['1200 press 1 Right 5','1400 press 1 Right 5','1600 press 1 Left 5','1800 press 1 Left 5','2000 exit'],checks);reports.append(result)

# All period lengths and all Line Changes choices, alternate goalie/fighting/difficulty.
source=out/'new.srm';seed(source,'legacy',0)
def first_game(case):
 period,line=case;v=[period,period%2,1+period%2,line,period%2,(period+line)%3];name=f'new-period{period}-lines{line}'
 result,battery=run(name,source,['1100 press 1 Start 5','1700 press 1 Start 5','2300 press 1 Start 5','3500 exit'],inject=v,capture=True)
 result2,_=run(name+'-reboot',battery,['1400 exit'],{1000:dict(slot=1,values=v)})
 return result,result2
with ThreadPoolExecutor(max_workers=3) as pool:
 for pair in pool.map(first_game,[(p,l) for p in range(4) for l in range(3)]):reports.extend(pair)
v=[3,1,2,2,1,2]
result,battery=run('real-first-game',source,['1100 press 1 Start 5','1700 press 1 Start 5','2300 press 1 Start 5','13001 exit'],inject=v,capture=True,finish=True);reports.append(result)
lines=[];checks={1000:dict(slot=1,values=v)};t=1100
for _ in range(4):lines.append(f'{t} press 1 Down 5');t+=90
for i in range(6):
 lines.append(f'{t} press 1 Right 5');checks[t+65]=dict(slot=1,values=v);t+=90
 if i<5:lines.append(f'{t} press 1 Down 5');t+=90
lines.append(f'{t+100} exit');result,_=run('real-first-game-cold-lock',battery,lines,checks);reports.append(result)
(out/'results.json').write_text(json.dumps(reports,indent=2)+'\n');print('PASS',len(reports),'settings lock/cold-boot/capture cases;',sum(r['checks'] for r in reports),'native observations')