from pathlib import Path
import json,sys,random,math
p=Path(outdir);spec=json.loads(os.environ['SHOT_CASE']);fixture=Path(spec['fixture']);meta=json.loads((fixture/'fixture.json').read_text());base=(fixture/'base.st').read_bytes();buf=C.create_string_buffer(base,len(base))
rb=lambda a:wram[ram(a)]
def rw(a):return rb(a)*256+rb(a+1)
def sw(a):v=rw(a);return v-65536 if v&32768 else v
def ww(a,v):v&=65535;wram[ram(a)]=v>>8;wram[ram(a+1)]=v&255
OBJ=lambda i:0xb04a+128*i
shooter=next(x['object'] for x in meta['actors'][:5] if x['slot']==spec['slot']);b=OBJ(shooter)
rng=random.Random(2601026);trials=[]
for i in range(spec['count']):
 dist=spec.get('distances',[180,200])[i%len(spec.get('distances',[180,200]))];x=[-40,-20,0,20,40][(i//2)%5]
 trials.append(dict(index=i+1,x=x,distance=dist,seed=rng.getrandbits(32)))
idx=-1;t=-1;c={};results=[];charge_at=spec.get('charge_at',4)
def snap(name):
 raw,w,h,pitch=frame[0]
 if pixfmt[0]==1:im=Image.frombuffer('RGBA',(w,h),raw,'raw','BGRA',pitch,1).convert('RGB')
 else:
  import struct
  px=bytearray()
  for yy in range(h):
   for (v,) in struct.iter_unpack('<H',raw[yy*pitch:yy*pitch+w*2]):px+=bytes(((v>>11)<<3,((v>>5)&63)<<2,(v&31)<<3))
  im=Image.frombytes('RGB',(w,h),bytes(px))
 im.save(p/(name+'.png'))
def test_tick(f):
 global idx,t,c
 if f<20:return
 if t<0:
  idx+=1
  if idx==len(trials):
   (p/'results.json').write_text(json.dumps(dict(spec=spec,fixture=meta,shots=results),indent=2));print('DONE',len(results),{r:sum(x['result']==r for x in results) for r in ['goal','save','over','wide','noshot']},'crits',sum(x.get('crit',False) for x in results),flush=True);sys.exit(0)
  assert core.retro_unserialize(buf,len(base));held.clear();t=0;c=dict(trials[idx]);c['score0']=rw(0xc6da);c['sog0']=rw(0xc6ce)
 tr=trials[idx];sgn=1 if rb(b+0x62)&128 else -1
 if t==0:
  for i in list(range(5))+list(range(6,11)):
   a=OBJ(i);j=i if i<5 else i-1;ww(a,-110+(j%5)*55);ww(a+2,0);ww(a+0x14,(-180+(j//5)*50)*sgn);ww(a+0x16,0);ww(a+0x28,0);ww(a+0x2a,0);wram[ram(a+0x62)]&=247
  wram[ram(b+0x62)]=(rb(b+0x62)&0xc0)|8;wram[ram(b+0x63)]=0;ww(0xc320,shooter)
  y=(261-tr['distance'])*sgn;ww(b,tr['x']);ww(b+2,0);ww(b+0x14,y);ww(b+0x16,0);ww(b+0x54,spec.get('facing',0))
  ww(0xb7aa,shooter);ww(0xb74a,tr['x']);ww(0xb75e,y+4*sgn)
  for a in [0xb74c,0xb760,0xb762,0xb772,0xb774,0xb776,0xf640,0xf000,0xf002,0xf0b0]:ww(a,0)
  ww(0xd066,tr['seed']>>16);ww(0xd068,tr['seed'])
  wram[ram(0xc2ea)]&=238;wram[ram(0xc2f2)]&=251
 if t>0:
  for i in list(range(5))+list(range(6,11)):
   if i==shooter:continue
   a=OBJ(i);j=i if i<5 else i-1;ww(a,-110+(j%5)*55);ww(a+0x14,(-180+(j//5)*50)*sgn);ww(a+0x28,0);ww(a+0x2a,0)
 if t==charge_at-2:
  ww(b+0x54,spec.get('facing',0));ww(b+0x28,0);ww(b+0x2a,0)
 if t==charge_at+37:
  ww(0xd066,tr['seed']>>16);ww(0xd068,tr['seed']);c['rng_release_input']=hex(tr['seed'])
 if t==charge_at:held[(0,IDS['C'])]=40
 if t>=charge_at+1:
  if rw(0xf646)>0 and 'release' not in c:
   c.update(release=t,rng_release_output=hex((rw(0xd066)<<16)|rw(0xd068)),crit=bool(rw(0xf060)),eligible=bool(rb(0xf0b0)),charge=rw(0xf644),animation=hex(rw(b+0x58)),distance_actual=rw(0xf000),bonus=bool(rb(0xf002)),power=rw(0xbede),velocity=[sw(0xb772),sw(0xb774)],vertical=sw(0xb776),kmh=rw(0xf650),release_xy=[sw(0xb74a),sw(0xb75e)],shooter_slot=rb(b+0x66),keeper_slot=rb(OBJ(11)+0x66),pct=[sw(0xf70e),sw(0xf72a)],ratings=bytes(rb(b+j) for j in range(0x66,0x76)).hex())
   if idx==0:snap('shot')
  if 'release' in c:
   if rw(0xc6da)>c['score0']:c['goal']=True
   if rw(0xc6ce)>c['sog0']:c['sog']=True
   px,py,ph=sw(0xb74a),sw(0xb75e)*sgn,sw(0xb762)
   if py>=261 and abs(px)<=24 and 'line' not in c:c['line']=[px,ph]
   owner=sw(0xb7aa)
   if owner>=0 and 'caught' not in c:c['caught']=dict(object=owner,frame=t)
  if ('release' in c and (c.get('goal') or c.get('caught') or t>=c['release']+140)) or t>=charge_at+240:
   c['debug_final']={'animation':hex(rw(b+0x58)),'anim_phase':rw(b+0x5a),'flags':hex(rw(b+0x62)),'owner':sw(0xb7aa),'control':rw(0xc320),'gameflags':[rb(0xc2ea),rb(0xc2ec),rb(0xc2f2)]};c['result']='noshot' if 'release' not in c else 'goal' if c.get('goal') else 'save' if c.get('sog') else 'over' if c.get('line',[0,0])[1]>13 else 'wide'
   results.append(c);(p/'shots.jsonl').open('a').write(json.dumps(c)+'\n');t=-1;return
 t+=1