import argparse, datetime, hashlib, json, pathlib, subprocess, time
from PIL import Image, ImageChops
p=argparse.ArgumentParser()
p.add_argument('--session', required=True)
p.add_argument('--out', required=True)
p.add_argument('--side', required=True)
p.add_argument('--sha', required=True)
p.add_argument('--capture-only', action='store_true')
a=p.parse_args()
out=pathlib.Path(a.out)/(a.side+'-'+datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%dT%H%M%SZ'))
out.mkdir(parents=True)
def call(*args):
    r=subprocess.run(['agent-device',*map(str,args),'--session',a.session],capture_output=True,text=True)
    if r.returncode: raise RuntimeError(r.stdout+r.stderr)
    return r.stdout
def snapshot(stage):
    data=json.loads(call('snapshot','--json'))['data']
    (out/(stage+'-raw.json')).write_text(json.dumps(data,indent=2))
    return data['nodes']
def press(label):
    nodes=snapshot('before-'+label.lower().replace(' ','-'))
    n=next(n for n in nodes if n.get('label')==label and n.get('kind')=='button' and n.get('hittable',True))
    r=n['rect']; call('press',r['x']+r['width']/2,r['y']+r['height']/2,'--settle')
if not a.capture_only:
    call('wait','text','New chat','30000')
    call('wait','stable','500','30000')
    press('New chat')
    call('wait','text','Choose project','30000')
    call('wait','stable','500','30000')
    press('VoiceOver demo')
    call('wait','text','New thread','30000')
    call('wait','stable','500','30000')
nodes=snapshot('composer')
composer=next(n for n in nodes if 'ComposerTextView' in n.get('role',''))
relevant=[n for n in nodes if 'ComposerTextView' in n.get('role','') or n.get('label')=='Add attachment' or n.get('label','').startswith('Model and reasoning')]
fields=('type','kind','role','label','value','editable','enabled','rect')
record={'commit':a.sha,'device':'iPhone 18 Pro','os':'iOS 27.2 simulator','viewport':{'width':402,'height':874},'theme':'light','screen':'Empty New thread','nodes':[{k:n.get(k) for k in fields} for n in relevant]}
(out/'accessibility.json').write_text(json.dumps(record,indent=2))
for attempt in range(3):
    first=out/f'still-{attempt}-a.png'; second=out/f'still-{attempt}-b.png'
    call('screenshot',first)
    time.sleep(.5)
    call('screenshot',second)
    with Image.open(first) as im1, Image.open(second) as im2:
        identical=im1.size==im2.size and ImageChops.difference(im1.convert('RGB'),im2.convert('RGB')).getbbox() is None
    if identical:
        (out/'composer.png').write_bytes(second.read_bytes())
        break
else: raise RuntimeError('No stable pair of screenshots')
(out/'stability.json').write_text(json.dumps({'identicalPixels':True,'minimumGapMs':500,'sha256':[hashlib.sha256(x.read_bytes()).hexdigest() for x in (first,second)]},indent=2))
print(json.dumps({'side':a.side,'label':composer.get('label'),'value':composer.get('value'),'output':str(out)}))