#!/usr/bin/env python3
"""Native Easy team skills, inverted awareness, identity and persistent fatigue.
Paired Normal/Hard controls and real input are recorded separately.
"""
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);a=ap.parse_args();h=Path(__file__).resolve().parent;o=a.output.resolve();o.mkdir(parents=True,exist_ok=True)
sha=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env['GPGX_DEV']='513'
def run(name,body,size,rom=None,edits=()):
 p=o/name;p.mkdir(exist_ok=True);b=[' .text',' .globl _start','_start:',' move.w #0x2700,%sr',' movea.l #0xFFF600,%sp',*body,' move.l #0x45415359,0xFF3FF0','end: bra.s end'];(p/'driver.s').write_text('\n'.join(b)+'\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();b=bytearray((rom or a.rom).read_bytes());assert len(code)<0x3000 and b[0x2d5000:0x2d5000+len(code)]==b'\xff'*len(code);b[0x2d5000:0x2d5000+len(code)]=code;struct.pack_into('>I',b,4,0x2d5000)
 for addr in [0x2e4004,0x2e5000,0x2ed000]:b[addr:addr+2]=bytes.fromhex('4e75')
 for addr,blob in edits:b[addr:addr+len(blob)]=blob
 (p/'driver.bin').write_bytes(b);(p/'commands.txt').write_text(f'500 dump ff3ff0 4 done\n500 dump ff4000 {size} result\n501 exit\n')
 with(p/'log').open('w') as f:subprocess.run(['nice','-n','19','python3',str(h/'test_core.py'),str(p/'driver.bin'),str(p/'commands.txt'),str(p)],env=env,stdout=f,stderr=f,check=True,timeout=60)
 r={l.split()[1]:bytes.fromhex(' '.join(l.split()[2:])) for l in (p/'commands.out').read_text().splitlines()};assert r['done']==b'EASY',name;return r['result']
# P1 through P4 alone, all co-op, head-to-head, and CPU-only; both rink directions.
layouts=[(0,0,0,0),(1,0,0,0),(0,2,0,0),(0,0,1,0),(0,0,0,2),(1,1,1,1),(2,2,2,2),(1,2,1,2)]
fields=[2,3,6,7,8,10,11,12,13,15]
def frame_case(spec):
 diff,layout,reverse=spec;tag=f'd{diff}-pads{"".join(map(str,layout))}-r{reverse}';b=[f' move.w #{diff},0xFFF07A',' move.w #0x4352,0xFFF700',' move.w #600,0xFFF702',' move.w #290,0xFFF70C',' move.w #290,0xFFC468']
 for off in [0x70e,0x72a,0x72c,0x72e,0x730,0x732,0x734]:b.append(f' clr.w 0xFFF{off:03X}')
 for pad,side in enumerate(layout):b.append(f' move.w #{side},0x{0xffc328+2*pad:X}')
 bases=[]
 for i in range(12):
  obj=0xffb04a+128*i;values=bytes((3+i+j)%31 for j in range(16));bases.append(values)
  b += [f' move.b #{(64 if i>=6 else 0)+(128 if bool(i>=6)^bool(reverse) else 0)},0x{obj+0x62:X}',f' move.b #1,0x{0xfff71c+i:X}',f' move.b #127,0x{0xfff710+i:X}']
  for j,v in enumerate(values):b += [f' move.b #{v},0x{obj+0x66+j:X}',f' move.b #{v},0x{0xfff740+16*i+j:X}']
 for team in [0xffc6ce,0xffca32]:
  for slot in range(26):b.append(f' move.w #{500+slot*100},0x{team+0x32+2*slot:X}')
 b += [' jsr 0x2E0000',' lea 0xFF4000,%a4']
 for i in range(12):
  for j in range(16):b.append(f' move.b 0x{0xffb04a+128*i+0x66+j:X},(%a4)+')
 for team in [0xffc6ce,0xffca32]:
  for slot in range(26):b.append(f' move.w 0x{team+0x32+2*slot:X},(%a4)+')
 current=run(tag,b,296)
 if diff!=2:
  parent=run(tag+'-parent',b,296,a.baseline);assert current==parent,(tag,'Normal/Hard changed')
 else:
  for i,base in enumerate(bases):
   got=current[i*16:(i+1)*16];human=(i//6+1) in layout;want=bytearray(base)
   for j in fields:want[j]=30 if human else 6
   want[4]=want[5]=7 if human else 13;want[14]=14 if human else 2
   assert got==want,(tag,i,got.hex(),want.hex())
  energies=struct.unpack('>52H',current[192:]);want=tuple(4096 if side in layout else 500+slot*100 for side in [1,2] for slot in range(26));assert energies==want,(tag,'full roster energy',energies,want)
 return dict(case=tag,all12_skills=True,all52_energies=True,paired_parent_unchanged=diff!=2)
def energy_case(spec):
 diff,side,human,slot,op,x=spec;team=[0xffc6ce,0xffca32][side-1];name=f'energy-{diff}-{side}-{human}-{slot}-{op}-{x}';body=[f' move.w #{diff},0xFFF07A']
 for pad in range(4):body.append(f' move.w #{side if pad==3 and human else 0},0x{0xffc328+2*pad:X}')
 body += [f' movea.l #0x{team:X},%a2',f' move.w #{slot*2},%d1',f' move.w #1234,0x{team+0x32+2*slot:X}',' move.l #0x1234007B,%d0',f' move.w #{0x2700+(16 if x else 0)},%ccr',f' jsr 0x{0x15774 if op=="read" else 0x15782:X}',' move.w %sr,0xFF4000',' movem.l %d0-%d7/%a0-%a6,0xFF4002',' move.l %sp,0xFF403E',f' move.w 0x{team+0x32+2*slot:X},0xFF4042']
 r=run(name,body,68);value=struct.unpack_from('>I',r,2)[0];energy=int.from_bytes(r[66:68],'big');want=1234 if op=='read' else 123
 assert energy==want,(name,energy,want)
 assert value==(0x12340000+(want if op=='read' else 123)),(name,hex(value))
 assert int.from_bytes(r[:2],'big')&0x1f==(16 if x else 0),(name,'CCR',r[:2].hex())
 assert int.from_bytes(r[62:66],'big')==0xfff600,(name,'stack')
 assert r==run(name+'-parent',body,68,a.baseline),(name,'native ABI differs')
 return dict(case=name,energy=energy,complete_ABI=True)
with ThreadPoolExecutor(max_workers=3) as pool:frames=list(pool.map(frame_case,[(d,l,r) for d in [0,1,2] for l in layouts for r in [0,1]]))
print('PASS',len(frames),'native whole-team/bench profiles, including paired Normal/Hard controls',flush=True)
with ThreadPoolExecutor(max_workers=3) as pool:energy=list(pool.map(energy_case,[(d,s,u,k,op,x) for d in [0,1,2] for s in [1,2] for u in [0,1] for k in [0,13,25] for op in ['read','write'] for x in [0,1]]))
print('PASS',len(energy),'native energy tail cases, CCR/stack and paired guards',flush=True)
(o/'results.json').write_text(json.dumps(dict(rom_sha256=sha(a.rom),baseline_sha256=sha(a.baseline),frames=frames,energy=energy),indent=2)+'\n')

# Cold boot through actual season menus, using only an initial SRAM configuration.
def input_case(count):
 p=o/f'input-{count}';p.mkdir(exist_ok=True);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([80,67,count,count^255,2,253]);b[0x1ff6:0x1ffe]=bytes([0x55,0,0,2,2,0,0x5a,0]);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')
 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;(p/'battery.srm').write_bytes(raw);(p/'spec.json').write_text(json.dumps(dict(count=count,end=12000)));(p/'commands.txt').write_text('1500 press 1 Start 5\n1900 press 1 Start 5\n9000 snap\n12001 exit\n')
 # A preserving observation of full roster energy at the committed Easy frame
 # boundary; VBlank may otherwise interrupt the native debit/refill sequence.
 code=(h/'testcases/easy_energy_probe.s').read_text();(p/'probe.s').write_text(code)
 subprocess.run([str(h.parent/'m68kasm.sh'),str(p/'probe.s'),'2D9000',str(p/'probe.bin')],check=True,stdout=subprocess.DEVNULL)
 data=bytearray(a.rom.read_bytes());probe=(p/'probe.bin').read_bytes();assert data[0x2d9000:0x2d9000+len(probe)]==b'\xff'*len(probe);data[0x2d9000:0x2d9000+len(probe)]=probe;at=data.find(bytes.fromhex('4eb9002ee400'),0x2e0000,0x2e0800);assert at>=0;data[at:at+6]=bytes.fromhex('4eb9002d9000');(p/'observed.bin').write_bytes(data)
 with(p/'log').open('w') as f:subprocess.run(['nice','-n','19','python3',str(h/'test_core.py'),str(p/'observed.bin'),str(p/'commands.txt'),str(p)],env=env|dict(GPGX_DEV='1281' if count==4 else '513',GPGX_SRM=str(p/'battery.srm'),GPGX_HOOK=str(h/'testcases/easy_input_hook.py')),stdout=f,stderr=f,check=True,timeout=120)
 return json.loads((p/'results.json').read_text())
with ThreadPoolExecutor(max_workers=2) as pool:inputs=list(pool.map(input_case,[1,4]))
r=json.loads((o/'results.json').read_text());r['input']=inputs;(o/'results.json').write_text(json.dumps(r,indent=2)+'\n');print('PASS cold-boot Easy real-input games, solo and P1-P4; all skills/keepers/full bench energy',flush=True)

# Execute the complete native loader twice (actual roster substitution), then
# native roster numeric/Overall evaluation and the active Easy crit lookup.
def loader_case(spec):
 side,human,keeper,reverse=spec;team=0xffc6ce+side*0x364;idx=side*6+(5 if keeper else 2);obj=0xb04a+idx*128;ptr=struct.unpack_from('>I',a.rom.read_bytes(),0x30e+4*10)[0];name=f'loader-{side}-{human}-{keeper}-{reverse}';body=[' lea 0xFFFFB000,%a0',' move.w #0x4800/4-1,%d0','clear_loader: clr.l (%a0)+',' dbra %d0,clear_loader',' move.w #2,0xFFF07A',' move.w #0x4352,0xFFF700',f' move.l #0x{ptr:X},0x{team+0x1e:X}',f' move.w #0x{0xb04a+side*768:X},0x{team+0x22:X}',f' move.w #{0 if keeper else 1},0x{0xff0000+obj+0x34:X}',f' move.w #{idx},0x{0xff0000+obj+0x52:X}',f' move.b #{64*side+(128 if reverse else 0)},0x{0xff0000+obj+0x62:X}',f' move.w #{side+1 if human else 0},0xFFC32E']
 for k in range(26):body.append(f' move.w #-1,0x{team+0x66+2*k:X}')
 slots=[0,1] if keeper else [13,25]
 for n,slot in enumerate(slots):
  body += [f' movea.w #0x{obj:X},%a3',f' move.w #{slot},%d3',' jsr 0x15AA4',f' lea 0x{0xff4000+16*n:X},%a0',f' lea 0x66(%a3),%a1',' moveq #15,%d0',f'copy{n}: move.b (%a1)+,(%a0)+',f' dbra %d0,copy{n}']
 body += [f' movea.w #0x{team&65535:X},%a2',' moveq #13,%d0',' move.l #0x4000000A,%d4',' jsr 0x2E000C',' move.w %d0,0xFF4020',' move.w %d1,0xFF4022',' swap %d4',' jsr 0x8E40',' move.w %d0,0xFF4026',f' movea.w #0x{obj:X},%a3',' moveq #10,%d0',' moveq #13,%d1',' jsr 0x2E0020',' move.w %d0,0xFF4024']
 raw=run(name,body,40,edits=[(0x8e4c,bytes.fromhex("4e75"))])
 for n,slot in enumerate(slots):
  row=raw[16*n:16*n+16];assert row[0]==slot,(name,'roster identity',row.hex());assert all(row[j]==(30 if human else 6) for j in fields) and row[4:6]==bytes([7 if human else 13]*2),(name,'native loaded skills',row.hex());assert row[14]==(14 if human else 2),(name,'fighting')
 assert struct.unpack_from('>3H',raw,32)==(100 if human else 20,100,100 if human else 20),(name,'display/crit',raw[32:].hex())
 assert int.from_bytes(raw[38:40],'big')==(100 if human else 20),(name,'actual native numeric formatter',raw[38:].hex())
 return dict(case=name,native_roster_substitutions=slots,loaded_ratings=raw[:32].hex(),roster_and_crit=struct.unpack_from('>3H',raw,32))
with ThreadPoolExecutor(max_workers=3) as pool:loaders=list(pool.map(loader_case,[(s,u,k,r) for s in [0,1] for u in [0,1] for k in [0,1] for r in [0,1]]))
r=json.loads((o/'results.json').read_text());r['native_loader_roster_crit']=loaders;(o/'results.json').write_text(json.dumps(r,indent=2)+'\n');print('PASS',len(loaders),'complete native loader/substitution/roster/crit profiles, keepers and skaters both ends',flush=True)