"""
Can you rely on the agent?  Steps 3 to 5.
=========================================
Data: tau-bench historical trajectories (Yao et al., 2024; github.com/sierra-research/tau-bench, MIT).
Customer-service agents (GPT-4o and Claude 3.5 Sonnet) talk to a simulated customer
and use tools to look up and change bookings (airline) or orders (retail), following
a written policy. Each task was run several times (4 for GPT-4o, 8 for Claude). A run
succeeds only if the database ends up exactly as the task requires.

Setup (once):  nothing to install beyond Python 3
Run:           python agents_lesson.py                      (success per run, every run, any run)
               python agents_lesson.py failures             (what went wrong in failed runs)
               python agents_lesson.py show airline gpt-4o 41 0
                                                            (read one run: domain, model, task, trial)
The first run downloads about 50 MB into ./data.
"""
import json
import os
import sys
import urllib.request
from collections import Counter, defaultdict
from math import comb

BASE = 'https://raw.githubusercontent.com/sierra-research/tau-bench/main/historical_trajectories/'
RUNS = {('airline', 'gpt-4o'): 'gpt-4o-airline', ('retail', 'gpt-4o'): 'gpt-4o-retail',
        ('airline', 'sonnet'): 'sonnet-35-new-airline', ('retail', 'sonnet'): 'sonnet-35-new-retail'}
WRITES = ('book', 'cancel', 'modify', 'update', 'exchange', 'return', 'send')   # tools that change the database


def load(domain, model):
    os.makedirs('data', exist_ok=True)
    name = RUNS[(domain, model)]
    path = f'data/tau_{name}.json'
    if not os.path.exists(path):
        print('downloading', name, '...')
        urllib.request.urlretrieve(BASE + name + '.json', path)
    return json.load(open(path, encoding='utf-8'))


def by_task(runs):
    tasks = defaultdict(list)
    for r in runs:
        tasks[r['task_id']].append(r['reward'] >= 0.999)
    return tasks


def all_k(tasks, k):
    """Chance that k runs of a task, picked at random, ALL succeed (often written pass^k)."""
    return sum(comb(sum(v), k) / comb(len(v), k) for v in tasks.values()) / len(tasks)


def any_k(tasks, k):
    """Chance that at least one of k runs succeeds (often written pass@k)."""
    return sum(1 - comb(len(v) - sum(v), k) / comb(len(v), k) for v in tasks.values()) / len(tasks)


def report():
    print(f"  {'domain':<8} {'model':<8} {'tasks':>5} {'runs each':>9} {'one run':>8} {'all 4 runs':>11} {'any of 4':>9}"
          f"   always / never / sometimes")
    for (domain, model) in RUNS:
        tasks = by_task(load(domain, model))
        v = list(tasks.values())
        always, never = sum(all(x) for x in v), sum(not any(x) for x in v)
        print(f"  {domain:<8} {model:<8} {len(v):>5} {len(v[0]):>9} {all_k(tasks, 1):>8.0%} {all_k(tasks, 4):>11.0%} {any_k(tasks, 4):>9.0%}"
              f"   {always:>6} / {never:>5} / {len(v) - always - never:>9}")


def tool_calls(run):
    out = []
    for m in run['traj']:
        for call in (m.get('tool_calls') or []):
            f = call.get('function', {})
            raw = f.get('arguments') or {}
            try:
                args = raw if isinstance(raw, dict) else json.loads(raw)   # the two models' logs store these differently
            except ValueError:
                args = {}
            out.append((f.get('name'), args))
    return out


def failure_kind(run):
    """A rough rule: compare the database-changing tool calls the agent made with the ones the task required."""
    need = [(a['name'], a['kwargs']) for a in run['info']['task']['actions'] if a['name'].startswith(WRITES)]
    made = [(n, a) for n, a in tool_calls(run) if n and n.startswith(WRITES)]
    need_n, made_n = Counter(n for n, _ in need), Counter(n for n, _ in made)
    if not made and need:
        return 'made no change'
    if need_n - made_n:
        return 'skipped a required change'
    if made_n - need_n:
        return 'made a change it should not have'
    key = lambda pairs: sorted(json.dumps(a, sort_keys=True) for _, a in pairs)
    if key(need) != key(made):
        return 'right kind of change, wrong details'
    return 'changes right, reply to the customer wrong'


def failures():
    print("Failed runs sorted by a rough rule (read some before trusting it)\n")
    for (domain, model) in RUNS:
        failed = [r for r in load(domain, model) if r['reward'] < 0.999]
        kinds = Counter(failure_kind(r) for r in failed)
        print(f"  {domain} / {model}: {len(failed)} failed runs")
        for kind, n in kinds.most_common():
            print(f"      {n / len(failed):>4.0%}  {kind}")


def show(domain, model, task_id, trial):
    run = [r for r in load(domain, model) if r['task_id'] == task_id and r['trial'] == trial][0]
    print("TASK (what the simulated customer was told):", run['info']['task']['instruction'])
    print("REQUIRED CHANGES:", [a['name'] for a in run['info']['task']['actions'] if a['name'].startswith(WRITES)] or 'none')
    print("RESULT:", "success" if run['reward'] >= 0.999 else "failure", "|", failure_kind(run) if run['reward'] < 0.999 else "")
    print()
    for m in run['traj'][1:]:
        if m['role'] == 'tool':
            print('     [tool result]', str(m['content'])[:300].replace('\n', ' '))
        elif m.get('tool_calls'):
            for n, a in tool_calls({'traj': [m]}):
                print(f'     [agent calls {n}]', json.dumps(a)[:200])
        else:
            print(f"{m['role'].upper()}: {str(m['content'])[:500]}")


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    if cmd == 'failures':
        failures()
    elif cmd == 'show' and len(sys.argv) == 6:
        show(sys.argv[2], sys.argv[3], int(sys.argv[4]), int(sys.argv[5]))
    else:
        report()