#!/usr/bin/env python3
"""Summarize tau3-bench banking_knowledge results per arm: pass, tokens by class, $ from usage.
usage: summarize.py <results.json> [...]   (prints one JSON line per file; per-task rows with --tasks)
Prices: claude-haiku-5-5 list (litellm model_cost): input $0.10/M, cache write $0.125/M, cache read $0.01/M,
output $0.50/M; above 100k prompt tokens: input $0.50/M, cache write $0.625/M, cache read $0.05/M, output $2.50/M."""
import json, sys, math
P = dict(i=0.10e-6, cw=0.125e-6, cr=0.01e-6, o=0.50e-6)
P_LONG = dict(i=0.50e-6, cw=0.625e-6, cr=0.05e-6, o=2.50e-6)

def msg_cost(u):
    pt = u.get('prompt_tokens') or 0
    cr = u.get('cache_read_input_tokens') or 0
    cw = u.get('cache_creation_input_tokens') or 0
    o = u.get('completion_tokens') or 0
    i = max(pt - cr - cw, 0)
    p = P_LONG if pt > 100_000 else P
    return dict(i=i, cw=cw, cr=cr, o=o, usd=i*p['i'] + cw*p['cw'] + cr*p['cr'] + o*p['o'], long=int(pt > 100_000))

def wilson(k, n, z=1.96):
    if n == 0: return (0, 0)
    ph = k/n; d = 1 + z*z/n; c = ph + z*z/(2*n); r = z*math.sqrt(ph*(1-ph)/n + z*z/(4*n*n))
    return ((c - r)/d, (c + r)/d)

def summarize(path, tasks=False):
    d = json.load(open(path))
    rows = []
    for s in d['simulations']:
        r = (s.get('reward_info') or {}).get('reward')
        acc = {'agent': dict(i=0, cw=0, cr=0, o=0, usd=0.0, calls=0, long=0), 'user': dict(i=0, cw=0, cr=0, o=0, usd=0.0, calls=0, long=0)}
        for m in s.get('messages') or []:
            role = m.get('role')
            u = (m.get('raw_data') or {}).get('usage')
            if not u: continue
            side = 'agent' if role == 'assistant' else 'user' if role == 'user' else None
            if not side: continue
            c = msg_cost(u)
            for k in ('i', 'cw', 'cr', 'o', 'usd', 'long'): acc[side][k] += c[k]
            acc[side]['calls'] += 1
        rows.append(dict(task=s['task_id'], reward=r, term=s.get('termination_reason'), agent=acc['agent'], user=acc['user']))
    n = len(rows); k = sum(1 for r in rows if (r['reward'] or 0) >= 1 - 1e-9)
    tot = lambda side, f: sum(r[side][f] for r in rows)
    out = dict(file=path, n=n, passed=k, pass_rate=k/n if n else 0, ci95=wilson(k, n),
               agent_usd=tot('agent', 'usd'), user_usd=tot('user', 'usd'),
               agent_tokens={f: tot('agent', f) for f in ('i', 'cw', 'cr', 'o')}, agent_calls=tot('agent', 'calls'),
               agent_long_calls=tot('agent', 'long'),
               user_tokens={f: tot('user', f) for f in ('i', 'cw', 'cr', 'o')})
    if tasks: out['tasks'] = rows
    return out

if __name__ == '__main__':
    t = '--tasks' in sys.argv
    for p in [a for a in sys.argv[1:] if not a.startswith('--')]:
        print(json.dumps(summarize(p, t)))