from pathlib import Path
import asyncio,json,sys,re,uuid,urllib.request,urllib.parse,datetime,os
import websockets
ROOT=Path('/tmp/lead-24568-r1/evidence-capture-2');RUN=Path((ROOT/'active-run.txt').read_text());SIDE=sys.argv[1];D=RUN/SIDE
now=lambda:datetime.datetime.now(datetime.timezone.utc).isoformat()
uid=lambda:str(uuid.uuid4())
def log(kind,**obj):
 with (D/'rpc.jsonl').open('a') as f:f.write(json.dumps({'timestamp':now(),'kind':kind,**obj})+'\n')
def connect_info():
 meta=json.loads((D/'server-process.json').read_text());text=Path(meta['log']).read_text();m=re.search(r'serverPort=(\d+) webPort=(\d+)',text);return int(m[1]),int(m[2]),text
async def main():
 port,web,text=connect_info();origin=f'http://localhost:{web}';backend=f'http://localhost:{port}'
 if (D/'private-auth.json').exists():auth=json.loads((D/'private-auth.json').read_text())
 else:
  url=re.search(r'pairingUrl: (\S+)',text)[1];parsed=urllib.parse.urlparse(url);credential=urllib.parse.parse_qs(parsed.fragment or parsed.query)['token'][0]
  req=urllib.request.Request(backend+'/api/auth/browser-session',data=json.dumps({'credential':credential}).encode(),headers={'Content-Type':'application/json','Origin':origin})
  with urllib.request.urlopen(req) as resp:
   data=json.load(resp);cookie='; '.join(x.split(';')[0] for x in resp.headers.get_all('Set-Cookie') if 'Max-Age=0' not in x)
  auth={'cookie':cookie,'response':data};(D/'private-auth.json').write_text(json.dumps(auth));os.chmod(D/'private-auth.json',0o600)
 async with websockets.connect(f'ws://localhost:{port}/ws?orchestrationProtocol=2',origin=origin,additional_headers={'Cookie':auth['cookie']},max_size=32*1024*1024) as ws:
  pending={};seq=0;events=asyncio.Queue()
  async def reader():
   async for raw in ws:
    frames=json.loads(raw);frames=frames if isinstance(frames,list) else [frames]
    for frame in frames:
     log('rpc-receive',frame=frame);events.put_nowait(frame)
     if frame.get('_tag')=='Ping':await ws.send(json.dumps({'_tag':'Pong'}));continue
     if frame.get('_tag')=='Chunk':await ws.send(json.dumps({'_tag':'Ack','requestId':frame['requestId']}))
     if frame.get('_tag')=='Exit' and str(frame.get('requestId')) in pending:pending[str(frame['requestId'])].set_result(frame['exit'])
  reader_task=asyncio.create_task(reader())
  async def call(tag,payload,stream=False):
   nonlocal seq
   seq+=1;i=str(seq);frame={'_tag':'Request','id':i,'tag':tag,'payload':payload,'headers':[]};log('rpc-send',frame=frame)
   if not stream:pending[i]=asyncio.get_running_loop().create_future()
   await ws.send(json.dumps(frame))
   if stream:return i
   result=await asyncio.wait_for(pending[i],45)
   if result.get('_tag')!='Success':raise RuntimeError(tag+' '+json.dumps(result))
   return result.get('value')
  async def command(payload):return await call('orchestration.dispatchCommand',{'commandId':uid(),**payload})
  action=sys.argv[2]
  ids=json.loads((D/'ids.json').read_text()) if (D/'ids.json').exists() else {}
  if action=='seed':
   provider={'driver':'codex','enabled':True,'displayName':'Evidence stub','environment':[{'name':'T3_CODEX_COLLAB_SCRIPT','value':str(D/'provider/script.json'),'sensitive':False}],'config':{'enabled':True,'binaryPath':str(D/'provider/codex'),'homePath':str(D/'provider/home')}}
   await call('server.updateSettings',{'patch':{'defaultAutoPull':False},'providerInstanceMutation':{'operation':'upsert','instanceId':'proof_codex','instance':provider}})
   ids={'project':uid(),'watch':uid(),'error':uid(),'hover':uid()};(D/'ids.json').write_text(json.dumps(ids))
   await call('projects.mutate',{'type':'project.create','commandId':uid(),'projectId':ids['project'],'title':'Azure watch proof','workspaceRoot':str(D/'project'),'defaultModelSelection':None})
   await call('projects.mutate',{'type':'project.update','commandId':uid(),'projectId':ids['project'],'autoPull':False})
   for key,title in [('watch','Azure build watch'),('error','Azure policy read error'),('hover','Azure cold hover')]:
    await command({'type':'thread.create','threadId':ids[key],'projectId':ids['project'],'createdBy':'user','creationSource':'web','title':title,'modelSelection':{'instanceId':'proof_codex','model':'gpt-5.4'},'runtimeMode':'full-access','interactionMode':'default','branch':'feat/page','worktreePath':None})
   await command({'type':'thread.pull-request.watch','threadId':ids['watch'],'host':'dev.azure.com','repository':'acme/platform/_git/web','number':42,'watching':True,'link':{'url':'https://dev.azure.com/acme/platform/_git/web/pullrequest/42','source':'manual'}})
   await command({'type':'message.dispatch','threadId':ids['hover'],'messageId':uid(),'text':'Preview fixture: https://dev.azure.com/acme/platform/_git/web/pullrequest/90','createdBy':'user','creationSource':'web','attachments':[],'dispatchMode':{'type':'start_immediately'}})
   print('seeded',SIDE,ids)
  elif action=='wait-idle':
   for key in ['watch','error']:
    await call('orchestration.subscribeThread',{'threadId':ids[key],'acceptBoundedSnapshot':True,'requestCompletionMarker':True},stream=True)
    p=await call('orchestration.getThreadProjection',{'threadId':ids[key]})
    while any(r['status'] not in ['completed','failed','cancelled','interrupted'] for r in p['runs']):
     frame=await asyncio.wait_for(events.get(),45)
     if any(v.get('event',{}).get('type')=='run.updated' for v in frame.get('values',[])):
      p=await call('orchestration.getThreadProjection',{'threadId':ids[key]})
   print('provider runs terminal',SIDE)
  elif action=='refresh':
   print((await call('server.refreshProviders',{'instanceId':'proof_codex','fresh':True}))['providers'][-1].get('version'))
  elif action=='project':
   key=sys.argv[3] if len(sys.argv)>3 else 'watch';p=await call('orchestration.getThreadProjection',{'threadId':ids[key]});
   for item in p.get('turnItems',[]):
    if item['type']=='notification':
     detail=await call('orchestration.getTurnItem',{'threadId':ids[key],'itemId':item['id']});(D/('notification-'+item['id'].replace(':','_').replace('%','_')+'.json')).write_text(json.dumps(detail,indent=2))
   phase=json.loads((D/'state.json').read_text())['phase'];stamp=datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%dT%H%M%SZ');(D/f'projection-{phase}-{key}-{stamp}.json').write_text(json.dumps(p,indent=2));print(json.dumps(p))
  elif action=='watch-error':
   await command({'type':'thread.pull-request.watch','threadId':ids['error'],'host':'dev.azure.com','repository':'acme/platform/_git/web','number':43,'watching':True,'link':{'url':'https://dev.azure.com/acme/platform/_git/web/pullrequest/43','source':'manual'}});print('watched error PR43')
  elif action=='unwatch-main':
   await command({'type':'thread.pull-request.watch','threadId':ids['watch'],'host':'dev.azure.com','repository':'acme/platform/_git/web','number':42,'watching':False});print('unwatched main PR42')
  elif action=='unwatch-all':
   for key,n in [('watch',42),('error',43)]:await command({'type':'thread.pull-request.watch','threadId':ids[key],'host':'dev.azure.com','repository':'acme/platform/_git/web','number':n,'watching':False})
  elif action=='endpoints':
   ref={'projectId':ids['project'],'host':'dev.azure.com','repository':'acme/platform/_git/web','number':44}
   await call('pullRequests.invalidate',{'reference':ref})
   for method in ['pullRequests.preview','pullRequests.detail']:
    try:print(method,json.dumps(await call(method,{**ref,**({'allowStale':False} if method.endswith('detail') else {})})))
    except RuntimeError as e:print(str(e))
  elif action=='seed-hover-cold':
   ids['hoverCold']=uid();(D/'ids.json').write_text(json.dumps(ids))
   await command({'type':'thread.create','threadId':ids['hoverCold'],'projectId':ids['project'],'createdBy':'user','creationSource':'web','title':'Azure cold hover','modelSelection':{'instanceId':'proof_codex','model':'gpt-5.4'},'runtimeMode':'full-access','interactionMode':'default','branch':'feat/page','worktreePath':None})
   await command({'type':'message.dispatch','threadId':ids['hoverCold'],'messageId':uid(),'text':'Preview fixture: https://dev.azure.com/acme/platform/_git/web/pullrequest/91','createdBy':'user','creationSource':'web','attachments':[],'dispatchMode':{'type':'start_immediately'}})
   print('seeded fresh cold hover',SIDE,ids['hoverCold'])
  elif action=='invalidate-hover':
   print(await call('pullRequests.invalidate',{'reference':{'projectId':ids['project'],'host':'dev.azure.com','repository':'acme/platform/_git/web','number':90}}))
  elif action=='pair':
   session_req=urllib.request.Request(backend+'/api/auth/session',headers={'Cookie':auth['cookie'],'Origin':origin})
   with urllib.request.urlopen(session_req) as resp:data=json.load(resp)
   csrf=data.get('csrfToken') or auth['response'].get('csrfToken')
   req=urllib.request.Request(backend+'/api/auth/pairing-token',data=json.dumps({}).encode(),headers={'Cookie':auth['cookie'],'Origin':origin,'Content-Type':'application/json',**({'x-t3-csrf-token':csrf} if csrf else {})})
   with urllib.request.urlopen(req) as resp:data=json.load(resp)
   (D/'private-browser-pair.json').write_text(json.dumps(data));print('pair saved')
  elif action=='observe':
   for key in ['watch','error','hover']:await call('orchestration.subscribeThread',{'threadId':ids[key],'acceptBoundedSnapshot':True,'requestCompletionMarker':True},stream=True)
   await reader_task
  else:raise RuntimeError(action)
  reader_task.cancel()
asyncio.run(main())
