#!/usr/bin/env python3
import argparse, datetime, json, pathlib, subprocess, time
from PIL import Image, ImageChops
ROOT=pathlib.Path('/tmp/lead-33496-r1/capture2')
p=argparse.ArgumentParser();p.add_argument('session');p.add_argument('run');p.add_argument('action');p.add_argument('args',nargs='*');a=p.parse_args()
out=ROOT/a.run;out.mkdir(exist_ok=True)
def command(*args):
    started=datetime.datetime.now(datetime.timezone.utc).isoformat()
    r=subprocess.run(['agent-device',*map(str,args),'--session',a.session],capture_output=True,text=True,cwd='/home/ubuntu/wt/pr-33496',timeout=180)
    with (out/'driver.jsonl').open('a') as f:f.write(json.dumps({'at':started,'args':list(map(str,args)),'exit':r.returncode,'output':r.stdout,'error':r.stderr})+'\n')
    print(r.stdout or r.stderr,flush=True)
    if r.returncode:raise RuntimeError(r.stderr or r.stdout)
    return r.stdout
def still(name):
    for attempt in range(8):
        first=out/f'{name}-settle-{attempt}-a.png';second=out/f'{name}-settle-{attempt}-b.png'
        flags=['--normalize-status-bar']+(['--pixel-density','3'] if a.session.endswith('-ios') else [])
        command('screenshot',first,*flags);time.sleep(.5);command('screenshot',second,*flags)
        if ImageChops.difference(Image.open(first).convert('RGB'),Image.open(second).convert('RGB')).getbbox() is None:
            target=out/f'{name}.png';target.write_bytes(second.read_bytes());print('SETTLED '+str(target));return
    raise RuntimeError('No pixel-identical full-screen pair for '+name)
if a.action=='still':still(a.args[0])
elif a.action=='start':command('record','start',out/(a.args[0]+'.mp4'),'--hide-touches','--quality','high')
elif a.action=='stop':command('record','stop')
elif a.action=='snapshot':
    (out/(a.args[0]+'.txt')).write_text(command('snapshot','-i'))
elif a.action=='sequence':
    still('expanded')
    command('gesture','pan',540,1150,0,-450,1200);still('A')
    command('gesture','pan',540,850,0,450,1200);still('B')
    command('gesture','pan',540,1750,0,-650,1400);still('C')
    command('gesture','fling','up',540,1800,1200);still('D')
elif a.action=='command':command(*a.args)