#!/usr/bin/env python3
"""Private native 3/4-player matches, season recording and two cold boots.

Fast cases position a winning OT puck across the actual goal line; the native
engine records goals, statistics and saves. --natural also plays two complete
Mini matches without changing the clock or scoring RAM. Installed saves and
controls are never touched. The seed supplies only existing preferences.
"""
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);ap.add_argument('--natural',action='store_true')
ap.add_argument('--warm-state',type=Path,help='optional Four-player Mini main-menu capture')
ap.add_argument('--warm-save',type=Path,help='matching private battery, with Light saved as One')
a=ap.parse_args();out=a.output.resolve();out.mkdir(parents=True,exist_ok=True)
here=Path(__file__).resolve().parent
def execute(p,hook,spec,lines):
    (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='1281',GPGX_DEV2='1281',GPGX_SRM=str(p/'battery.srm'),GPGX_HOOK=str(here/hook),LIFECYCLE=json.dumps(spec),REBOOT_CASE=json.dumps(spec))
    with (p/'log.txt').open('w') as log:
        subprocess.run(['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=300)
    return json.loads((p/'result.json').read_text())
def run(case):
    slot,count,fast=case;length={1:84,2:31,3:7}[slot]
    p=out/f"slot{slot}-pads{count}-{'fast' if fast else 'natural'}"
    p.mkdir(exist_ok=True)
    b=bytearray(a.seed.read_bytes()[1:16384:2]);b[0xe20:0x1ef6]=bytes(0x10d6)
    struct.pack_into('>4s6H',b,0x1000,b'SESN',0,0,0,length,slot,3)
    b[0x1f08:0x1f12]=bytes(10);b[0x1fff]=sum(b[:0x1ffe])%256;b[0x1ffe]=b[0x1fff]^255
    raw=bytearray(16384);raw[1::2]=b;(p/'battery.srm').write_bytes(raw)
    lines=['1000 press 1 Down 5','1100 press 1 Right 5','1200 press 1 Right 5']
    if count==4:lines+=['1300 press 1 Right 5']
    lines+=['1400 press 1 Up 5','1500 press 1 Start 5','2100 press 1 C 5','180001 exit']
    result=execute(p,'four_lifecycle_hook.py',dict(slot=slot,count=count,length=length,fast=fast,run=1),lines)
    assert result['games']==2 and result['header'][2]==2
    # Start a fresh core/process, reselect the extra pads (TM is session RAM).
    boot=out/(p.name+'-reboot');boot.mkdir(exist_ok=True);(boot/'battery.srm').write_bytes((p/'battery.srm').read_bytes())
    # Player counts persist since r20. Resume the actual last game's stored
    # count; cycling a hard-coded number of times silently selects another mode.
    saved=(boot/'battery.srm').read_bytes()[1::2]
    stored_count=saved[0x1ef2]
    assert saved[0x1ef0:0x1ef2]==b'PC' and stored_count in (3,4)
    lines=['1500 press 1 Start 5','2100 press 1 C 5','12001 exit']
    execute(boot,'four_pad_reboot_hook.py',dict(slot=slot,length=length,count=stored_count,stage=1),lines)
    # Close an unfinished match and cold boot again. Progress must stay at 2.
    again=out/(p.name+'-restart');again.mkdir(exist_ok=True);(again/'battery.srm').write_bytes((boot/'battery.srm').read_bytes())
    lines=[];checks=[]
    for i in range(12):
        t=1450+150*i;lines.append(f'{t} press 1 {"Right" if i<6 else "Left"} 5');checks.append(t+110)
    lines+=['3500 snap','3501 exit']
    execute(again,'four_pad_reboot_hook.py',dict(slot=slot,length=length,count=count,stage=2,checks=checks,finish=3500),lines)
    return dict(slot=slot,count=count,fast=fast,result=result,reboots=2)
cases=[(slot,count,True) for slot in [1,2,3] for count in [3,4]]
if a.natural:cases.append((3,4,False))
with ThreadPoolExecutor(max_workers=2) as pool:results=list(pool.map(run,cases))
if a.warm_state:
    assert a.warm_save
    p=out/'warm-switch-one';p.mkdir(exist_ok=True);(p/'battery.srm').write_bytes(a.warm_save.read_bytes())
    hook=p/'hook.py';hook.write_text('''from pathlib import Path
def rw(a):return (wram[ram(a)]<<8)|wram[ram(a+1)]
def test_tick(f):
 if f==350:assert rw(0xd04a)==1 and rw(0xf07c)==4,('not the stale TM fixture',rw(0xd04a),rw(0xf07c))
 if f==2000:
  assert [rw(0xc328+2*i) for i in range(4)]==[1,0,0,0],('One started extra humans',[rw(0xc328+2*i) for i in range(4)])
  Path(outdir,'result.json').write_text('{"visible_one_started_one":true}');sys.exit(0)
''')
    execute(p,hook,{},[f'100 load {a.warm_state.resolve()}','200 press 1 Left 5','400 press 1 Start 5','1000 press 1 C 5','2001 exit'])
(out/'results.json').write_text(json.dumps(results,indent=2)+'\n')
print('PASS',len(results),'two-match lifecycles, two cold boots each; extra-pad movement/pause/resume/menu circles')