#!/usr/bin/env python3
"""Native smart-target arithmetic/guards/height, C30C displaced stack and ABI."""
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('--part',choices=['policy','hook','all'],default='all')
ap.add_argument('--benchmark',action='store_true')
a=ap.parse_args();h=Path(__file__).resolve().parent;o=a.output.resolve();o.mkdir(parents=True,exist_ok=True)
from smart_aim_test_support import targets, expected, setup

assert a.rom.read_bytes()[0xc542:0xc566]==struct.pack('>18h',*[v for t in targets for v in t]),'native target table changed'


def execute(s,p,kind,baseline=False):
    p.mkdir(parents=True,exist_ok=True);lines,addresses=setup(s)
    body=[' .text',' .globl _start','_start:',' move.w #0x2700,%sr',' movea.l #0xfffff900,%sp']+lines
    if s.get('friendly_interception'):
        slot=6*s.get('team',0)+2
        body += [' move.w #0x5042,0xfffff086',f' move.w #{slot+1},0xfffff088',' move.w #60,0xfffff08a',f' move.w #{0xFF00|slot},0xffffbee0',f' move.w #{6*s.get("team",0)+1},0xffffbeda']
    if kind=='policy':body += [' jsr 0x2df008',' move.w %sr,0xffff4800',' movem.l %d0-%d7/%a0-%a6,0xffff4000',' move.l %sp,0xffff403c']
    else:body += [' jsr 0xc30c']
    body += [' move.w 0xffffbedc,0xffff4802',' move.l 0xffffd066,0xffff480c',' move.l #0x5341494d,0xffff3ff0','end: bra.s end','observe:',' move.w %sr,0xffff4800',' movem.l %d0-%d7/%a0-%a6,0xffff4000',' move.w (%sp),0xffff4804',' move.l %sp,0xffff403c',' addq.l #2,%sp',' rts']
    (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')],stdout=subprocess.DEVNULL,check=True)
    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)
    if kind=='hook':
        matches=[code.index(bytes.fromhex(m)) for m in ['40f9ffff4800','40f84800'] if bytes.fromhex(m) in code];assert len(matches)==1;at=matches[0];data[0xc314:0xc31a]=bytes.fromhex('4ef9')+struct.pack('>I',0x2d9000+at)
        if baseline:data[0xc30c:0xc314]=bytes.fromhex('610002583f3c0005')
    (p/'driver.bin').write_bytes(data);(p/'commands.txt').write_text('100 dump ff3ff0 4 done\n100 dump ff4000 64 regs\n100 dump ff4800 16 observed\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'SAIM',s
    return out,addresses

rows=[]
if a.part in ('policy','all'):
    cases=[dict(team=t,direction=d,gx=g,index=i) for t in (0,1) for d in (-1,1) for g in (-20,20) for i in range(9)]
    cases += [dict(**s) for s in [dict(gx=0),dict(gx=7),dict(px=-48,depth=240,gx=-32),dict(depth=87),dict(depth=264),dict(depth=270),dict(gy=210),dict(gy=265),dict(gvx=-384,gvy=512),dict(keeper_slot=4),dict(keeper_role=-1),dict(keeper_boxed=True),dict(busy=True,animation=0x7fc),dict(busy=True,animation=0x92e)]]
    cases += [dict(**{key:value}) for key,value in [('difficulty',0),('difficulty',2),('human',True),('boxed',True),('role',0),('role',-1),('busy',True),('whistle',True),('highlight',True),('view',True),('offside',True),('penalty',True),('special',1),('special',4),('wrong_owner',True),('loose',True)]]
    cases += [dict(automatic_replay=v) for v in (0xa500,0xa501,0xa502,0xa506)]
    cases += [dict(team=t,direction=d,busy=True,animation=0xB24,friendly_interception=True)for t in(0,1)for d in(-1,1)]
    if a.benchmark:
        cases=[]
        for team in (0,1):
            for direction in (-1,1):
                for hard in (False,True):
                    hc=10 if (team==0)==hard else 5
                    cases.append(dict(team=team,direction=direction,difficulty=0,effective_difficulty=1 if hard else 0,home_club=hc,away_club=15-hc))
    for i,s in enumerate(cases):
        out,addresses=execute(s,o/('policy-'+str(i)),'policy');regs=struct.unpack('>16I',out['regs']);want=expected(s,s.get('index',0));assert (regs[0],regs[1],regs[2]&65535)==want,(s,regs[:3],want)
        assert list(regs[3:8])==[0x11223300+r for r in range(3,8)] and list(regs[8:15])==addresses and regs[15]==0xfffff900,(s,'register/stack ABI',regs)
        assert out['observed'][2:4]==struct.pack('>H',s.get('index',0)) and out['observed'][12:16]==bytes.fromhex('12345678'),(s,'aim or RNG mutated')
        assert targets[want[1]][1]==targets[s.get('index',0)][1]
        rows.append(dict(part='policy',case=s,active=want[0],index=want[1],gap=want[2],native_height_and_register_stack_RNG_preserved=True))
    print('PASS',len(cases),'native target/height/guard/bounds/ABI cases',flush=True)
if a.part in ('hook','all'):
    cases=[dict(team=t,direction=d,difficulty=level) for t in (0,1) for d in (-1,1) for level in (0,1,2)]
    cases += [dict(team=t,direction=d,human=True,index=4) for t in (0,1) for d in (-1,1)]
    for i,s in enumerate(cases):
        base,_=execute(s,o/('hook-old-'+str(i)),'hook',True);out,_=execute(s,o/('hook-new-'+str(i)),'hook');oldregs=list(struct.unpack('>16I',base['regs']));newregs=list(struct.unpack('>16I',out['regs']));original=int.from_bytes(base['observed'][2:4],'big');active,index,gap=expected(s,original)
        want=oldregs[:]
        if active and index!=original:want[0]=(want[0]&0xffff0000)|index
        assert newregs==want,(s,newregs,want)
        assert out['observed'][0:2]==base['observed'][0:2],(s,'native CCR mismatch')
        assert out['observed'][4:6]==base['observed'][4:6]==bytes.fromhex('0005') and newregs[15]==0xfffff8fa,(s,'displaced native word/stack')
        assert int.from_bytes(out['observed'][2:4],'big')==(index if active else original)
        assert out['observed'][12:16]==base['observed'][12:16]==bytes.fromhex('12345678')
        rows.append(dict(part='C30C_native_hook',case=s,native_index=original,selected_index=index if active else original,displaced_word5_stack_CCR_native_register_results_preserved=True))
    print('PASS',len(cases),'actual native C566/C30C displaced-stack and exact Normal/Easy/human contracts',flush=True)
(o/'results.json').write_text(json.dumps(dict(rom_sha256=hashlib.sha256(a.rom.read_bytes()).hexdigest(),cases=rows),indent=2)+'\n')