#!/usr/bin/env python3
"""Native taken recovery branch/hold/guard/ABI and actual carried-puck target contracts."""
from pathlib import Path
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)
a=ap.parse_args();here=Path(__file__).resolve().parent;tools=here.parent
o=a.output.resolve();o.mkdir(parents=True,exist_ok=True)
cases=[dict(team=t,sign=sgn,anim=anim,active=True) for t in (0,1) for sgn in (-1,1) for anim in (0x250,0x2a2)]
for key,value in [('duration',8),('duration',7),('duration',0),('duration',-1),('role',4),('human',8),('human',32),('anim',2),('anim',0x2f4),('phase',0),('owner',-1),('owner',-248),('owner',12),('own_team',True),('receiver_role',0),('receiver_boxed',4),('whistle',1),('whistle',16),('highlight',128),('special',1),('special',4),('keeper_x',0),('velocity',-511),('flags',0)]:
 cases.append(dict(active=False,**{key:value}))
cases += [dict(active=True,velocity=-512),dict(active=True,velocity=-513),dict(active=True,duration=9),dict(active=True,duration=64)]
cases += [dict(active=True,sprite_delay=v) for v in (0,1,127,128,255)]

def native(profile,rom,ss,batch):
 p=o/f'{profile}-{batch}';p.mkdir(exist_ok=True)
 body=[' .text',' .globl _start','_start:',' move.w #0x2700,%sr',' movea.l #0xFFF800,%sp']
 expected=[]
 for n,s in enumerate(ss):
  team=s.get('team',0);sgn=s.get('sign',1);g=6*team+5;own=s.get('owner',6*(team if s.get('own_team') else 1-team)+2);gb=0xffb04a+128*g;cb=0xffb04a+128*(own if 0<=own<12 else 2)
  actor=bytearray(128)
  vals=[(0,s.get('keeper_x',16)*sgn),(0x14,246),(0x34,s.get('role',0)),(0x52,g),(0x58,s.get('anim',0x250)),(0x5a,s.get('phase',4)),(0x5c,s.get('duration',32))]
  body += [f' lea 0x{gb:X},%a0',' moveq #31,%d0',f'clear{n}: clr.l (%a0)+',f' dbra %d0,clear{n}']
  for off,v in vals:
   body.append(f' move.w #{v},0x{gb+off:X}');struct.pack_into('>H',actor,off,v&65535)
  byvals=[(0x40,5),(0x62,team*64+s.get('human',0)),(0x63,s.get('flags',2)),(0x65,s.get('sprite_delay',100))]
  for off,v in byvals:body.append(f' move.b #{v},0x{gb+off:X}');actor[off]=v
  body += [f' move.w #{s.get("receiver_role",4)},0x{cb+0x34:X}',f' move.w #{s.get("velocity",-1200)*sgn},0x{cb+0x28:X}',f' move.b #{(own//6)*64 if 0<=own<12 else 64},0x{cb+0x62:X}',f' move.b #{s.get("receiver_boxed",0)},0x{cb+0x63:X}',f' move.w #{own},0xFFB7AA',f' move.b #{s.get("whistle",0)},0xFFC2EA',f' move.b #{s.get("highlight",0)},0xFFC2EC',f' move.b #{s.get("special",0)},0xFFC2FA']
  regs=[0x11223300+i for i in range(8)];regs[0]=(regs[0]&0xffff0000)|(s.get('duration',32)&65535);regs[7]=0xaabbcc01
  for i,v in enumerate(regs):body.append(f' move.l #0x{v:X},%d{i}')
  for i in range(7):body.append(f' movea.l #0x{gb if i==3 else 0xff1000+256*i:X},%a{i}')
  at=0xff4000+224*n
  body += [' move.w #0x271F,%sr',' jsr 0xD648',f' movem.l %d0-%d7/%a0-%a6,0x{at:X}',f' move.w %sr,0x{at+64:X}',f' move.l %sp,0x{at+68:X}',f' lea 0x{gb:X},%a0',f' lea 0x{at+80:X},%a1',' moveq #31,%d0',f'copy{n}: move.l (%a0)+,(%a1)+',f' dbra %d0,copy{n}']
  active=profile=='candidate' and s['active'];struct.pack_into('>H',actor,0x5c,(8 if active else s.get('duration',32))&65535);actor[0x65]=s.get('sprite_delay',100)
  if active:actor[0x40]=0
  ccr=0x1b if s.get('flags',2)&2 else 0x1f
  expected.append((regs+[gb if i==3 else 0xff1000+256*i for i in range(7)],0x2700|ccr,bytes(actor)))
 body += [' move.l #0x474C4154,0xFF3FF0','end: bra.s end']
 (p/'driver.s').write_text('\n'.join(body)+'\n')
 subprocess.run([str(tools/'m68kasm.sh'),str(p/'driver.s'),'2D9000',str(p/'code.bin')],check=True,stdout=subprocess.DEVNULL)
 r=bytearray(rom.read_bytes());code=(p/'code.bin').read_bytes();assert len(code)<0x3000 and r[0x2d9000:0x2d9000+len(code)]==b'\xff'*len(code)
 r[0x2d9000:0x2d9000+len(code)]=code;struct.pack_into('>I',r,4,0x2d9000);r[0xd652:0xd654]=bytes.fromhex('4e75');(p/'driver.bin').write_bytes(r)
 (p/'commands.txt').write_text(f'500 dump ff3ff0 4 done\n500 dump ff4000 {224*len(ss)} result\n501 exit\n')
 env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env['GPGX_DEV']='513'
 with (p/'log.txt').open('w') as log:subprocess.run(['python3',str(here/'test_core.py'),str(p/'driver.bin'),str(p/'commands.txt'),str(p)],env=env,stdout=log,stderr=log,check=True,timeout=30)
 data={l.split()[1]:bytes.fromhex(' '.join(l.split()[2:])) for l in (p/'commands.out').read_text().splitlines()};assert data['done']==b'GLAT'
 rows=[]
 for n,(s,(regs,sr,actor)) in enumerate(zip(ss,expected)):
  raw=data['result'][n*224:(n+1)*224]
  assert list(struct.unpack_from('>15I',raw))==regs,('registers',profile,s)
  assert struct.unpack_from('>H',raw,64)[0]==sr,('CCR',profile,s,raw[64:66].hex(),hex(sr))
  assert struct.unpack_from('>I',raw,68)[0]==0xfff800,('stack',profile,s)
  assert raw[80:208]==actor,('object fields',profile,s,raw[80:208].hex(),actor.hex())
  rows.append(dict(profile=profile,spec=s,timer=struct.unpack_from('>H',actor,0x5c)[0],all_registers_stack_CCR_and_other_object_fields_exact=True))
 return rows

rows=[]
for profile,rom in [('baseline',a.baseline),('candidate',a.rom)]:
 for batch in range((len(cases)+7)//8):rows += native(profile,rom,cases[batch*8:batch*8+8],batch)
# Only the carried-puck target's two-byte ASR->NOP is allowed in this native span.
before=a.baseline.read_bytes();after=a.rom.read_bytes()
assert before[0xd8e2:0xd902]!=after[0xd8e2:0xd902]
assert before[0xd8e2:0xd8f8]==after[0xd8e2:0xd8f8] and before[0xd8fa:0xd902]==after[0xd8fa:0xd902]
assert before[0xd8f8:0xd8fa]==bytes.fromhex('e240') and after[0xd8f8:0xd8fa]==bytes.fromhex('4e71')
(o/'results.json').write_text(json.dumps(dict(rom_sha256=hashlib.sha256(after).hexdigest(),baseline_sha256=hashlib.sha256(before).hexdigest(),cases=rows,guarded_native_target_span=True),indent=2)+'\n')
print('PASS',len(rows),'native duration/guard/whole-object/register/stack/CCR cases and guarded carry-target span')