#!/usr/bin/env python3
import subprocess,sys,os,threading,json,shutil
from pathlib import Path
root=Path(__file__).parent
os.environ['GROK_HOME']=str(root/'grok-home')
log=open(root/'acp.jsonl','a',buffering=1)
log.write(json.dumps({'argv':sys.argv[1:]})+'\n')
p=subprocess.Popen([shutil.which('grok'),*sys.argv[1:]],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=open(root/'grok-stderr.log','a'))
requests={}
auth_attempted=False
pending=None
def pump():
 for line in sys.stdin.buffer:
  try: log.write(json.dumps({'direction':'client','message':json.loads(line)})+'\n')
  except: pass
  m=json.loads(line)
  if 'id' in m: requests[m['id']]=m
  p.stdin.write(line);p.stdin.flush()
threading.Thread(target=pump,daemon=True).start()
for line in p.stdout:
 try:
  m=json.loads(line)
  if m.get('error',{}).get('code')==-32000 and not auth_attempted:
   auth_attempted=True;pending=requests.get(m['id'])
   p.stdin.write((json.dumps({'jsonrpc':'2.0','id':999999,'method':'authenticate','params':{'methodId':'xai.api_key'}})+'\n').encode());p.stdin.flush();continue
  if m.get('id')==999999:
   log.write(json.dumps({'authShimResponse':m})+'\n')
   p.stdin.write((json.dumps(pending)+'\n').encode());p.stdin.flush();continue
  if isinstance(m.get('error',{}).get('data'),str):
   m['error'].pop('data');line=(json.dumps(m)+'\n').encode()
 except: pass
 try: log.write(json.dumps({'direction':'grok','message':json.loads(line)})+'\n')
 except: pass
 sys.stdout.buffer.write(line);sys.stdout.buffer.flush()