import json, pathlib, sqlite3, sys, time
home, out = map(pathlib.Path, sys.argv[1:])
seen = set()
for side in ['base', 'head']:
    deadline = time.monotonic() + 600
    while time.monotonic() < deadline:
        try:
            db = sqlite3.connect(f'file:{home}/userdata/statev2.sqlite?mode=ro', uri=True)
            runs = [json.loads(x[0]) for x in db.execute('select payload_json from orchestration_v2_projection_runs order by ordinal')]
            attempts = [json.loads(x[0]) for x in db.execute('select payload_json from orchestration_v2_projection_run_attempts')]
            turns = [json.loads(x[0]) for x in db.execute('select payload_json from orchestration_v2_projection_provider_turns')]
            db.close()
            if len(runs) < 3 or runs[0]['threadId'] in seen or runs[2]['status'] not in ['completed','failed']:
                time.sleep(.25); continue
            seen.add(runs[0]['threadId'])
            by_attempt = {a['id']: a for a in attempts}
            by_run = {r['id']: r for r in runs}
            observations = {
                'side': side,
                'runs': [{'ordinal':r['ordinal'],'status':r['status'], 'attempts':[{'ordinal':a['attemptOrdinal'],'reason':a['reason'],'status':a['status']} for a in attempts if a['runId']==r['id']]} for r in runs],
                'providerTurns': [{'runOrdinal': by_run[by_attempt[t['runAttemptId']]['runId']]['ordinal'], 'attemptOrdinal':by_attempt[t['runAttemptId']]['attemptOrdinal'], 'status':t['status'],'nodeMatchesCurrentRunRoot':t['nodeId']==by_run[by_attempt[t['runAttemptId']]['runId']]['rootNodeId']} for t in turns],
            }
            log = [json.loads(line) for line in (home/'agent-log.jsonl').read_text().splitlines()]
            prompt_sessions = [x['sessionId'] for x in log if x.get('method')=='session/prompt']
            observations['agent']={'promptCount':len(prompt_sessions),'promptSessionCount':len(set(prompt_sessions)),'cancelCount':sum(x.get('method')=='session/cancel' for x in log)}
            (out/f'{side}-runtime.json').write_text(json.dumps(observations,indent=2)+'\n')
            print(side, json.dumps(observations), flush=True)
            break
        except (sqlite3.Error, FileNotFoundError, KeyError, json.JSONDecodeError):
            pass
        time.sleep(.25)
    else:
        raise RuntimeError('No completed observation for '+side)