#!/usr/bin/env python3
"""Commands used for the manual Grok restart proof. Synthetic data only.
Run with Python 3.12. Use each checkout's own fresh .t3 home and an isolated,
authenticated GROK_HOME. Start its dev server with T3CODE_DEV_SEED=0 and the
same GROK_HOME; enable Grok and custom model grok-4.5 in Settings.

First: project-add, start. Stop/restart the dev server, then: send THREAD_ID.
Open Usage > Tokens > Past 24h; filter to Grok Build, refresh, and capture.
Next: copy-events, refresh, and confirm the total stays the same.
For upgrade: stop head server, run upgrade-cache with base's v16 cache and
head's userdata directory. Run head with the unchanged base GROK_HOME, refresh
Usage and capture. This uses base-generated cache rows without fabricating them.
"""
import argparse, json, os, pathlib, shutil, subprocess

p = argparse.ArgumentParser(description=__doc__)
p.add_argument('operation', choices=['project-add','start','send','copy-events','summarize','upgrade-cache'])
p.add_argument('--checkout', type=pathlib.Path, default=pathlib.Path.cwd())
p.add_argument('--project', type=pathlib.Path)
p.add_argument('--home', type=pathlib.Path, help='Isolated GROK_HOME')
p.add_argument('--thread')
p.add_argument('--cache', type=pathlib.Path)
p.add_argument('--destination', type=pathlib.Path)
a = p.parse_args()

def cli(*args):
    return subprocess.run(['node','apps/server/src/bin.ts',*args,'--base-dir','.t3'],cwd=a.checkout,check=True)

def events():
    for f in (a.home/'sessions').rglob('updates.jsonl'):
        for line in f.open():
            d=json.loads(line)
            if d.get('params',{}).get('update',{}).get('sessionUpdate')=='turn_completed':
                yield d

if a.operation == 'project-add':
    a.project.mkdir(parents=True,exist_ok=True)
    (a.project/'README.md').write_text('Synthetic Grok usage demonstration.\n')
    subprocess.run(['git','init','-q',str(a.project)],check=True)
    cli('project','add',str(a.project),'--title','Grok restart proof')
elif a.operation == 'start':
    cli('thread','start',str(a.project),'Reply with just the word one.','--title','Resumed Grok turns','--provider','grok','--model','grok-4.5','--wait','--json')
elif a.operation == 'send':
    cli('thread','send',a.thread,'Reply with just the word two.','--wait','--json')
elif a.operation == 'copy-events':
    src=next((a.home/'sessions').rglob('updates.jsonl'))
    dst=a.home/'sessions'/'exact-copy'/'updates.jsonl'
    dst.parent.mkdir(exist_ok=True)
    shutil.copyfile(src,dst)
elif a.operation == 'summarize':
    unique={}
    for e in events():
        q=e['params'];u=q['update'];key=(q['sessionId'],u['prompt_id'],q['_meta']['agentTimestampMs'])
        unique[key]=u['usage']
    print(json.dumps({'distinctTurns':len(unique),'processedTokens':sum(u['totalTokens'] for u in unique.values()),'reportedCostUsd':sum(u['costUsdTicks'] for u in unique.values())/1e10,'turns':[{'sessionId':k[0],'promptId':k[1],'timestampMs':k[2],'usage':u} for k,u in unique.items()]},indent=2))
else:
    d=json.loads(a.cache.read_text())
    files={k:v for k,v in d['files'].items() if v['p']=='grok'}
    models=[];sessions=[]
    for v in files.values():
        for r in v['r']+v['t']:
            m=d['models'][r[1]];s=d['sessions'][r[2]]
            if m not in models: models.append(m)
            if s not in sessions: sessions.append(s)
            r[1]=models.index(m);r[2]=sessions.index(s)
            if len(r)>11 and r[11] is not None:
                m=d['models'][r[11]]
                if m not in models: models.append(m)
                r[11]=models.index(m)
    out={'version':16,'models':models,'sessions':sessions,'files':files}
    a.destination.mkdir(parents=True,exist_ok=True)
    (a.destination/'usage-scan-cache-v16.json').write_text(json.dumps(out))
    (a.destination/'usage-scan-cache-v17.json').unlink(missing_ok=True)