#!/usr/bin/env python3
"""Agents Finance template based on cfo-bench (github.com/tom-eberle/cfo-bench, MIT).

Reads cfo-bench's public results/results.jsonl and writes three CSV tables that the
cfo-bench README does not publish in this shape:

  by_task.csv     one row per request (36): tries done right per agent (out of 3)
  by_family.csv   one row per agent and request type: success, safe tries and
                  API-equivalent cost per request done right
  unsafe.csv      one row per unsafe try: task, agent and the kinds of harm the
                  graders recorded

It also re-checks the headline numbers (success and safe tries per agent) so you can
compare them with results/summary.md. Only scored rows (status == "scored") count for
success and safety; replaced attempts are added to cost, as cfo-bench does.

Usage:  python3 cfobench_breakdown.py path/to/results.jsonl [--out DIR]
Licence: MIT-0 (Agents Finance). No network, standard library only.
"""
import argparse
import csv
import json
import os
from collections import defaultdict

AGENTS = ["claude-opus", "codex", "claude-sonnet", "claude-haiku", "codex-luna"]


def load(path):
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def pct(n, d):
    return f"{100 * n / d:.0f}" if d else ""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("results")
    ap.add_argument("--out", default="out")
    args = ap.parse_args()
    rows = load(args.results)
    os.makedirs(args.out, exist_ok=True)
    scored = [r for r in rows if r.get("status") == "scored"]
    agents = [a for a in AGENTS if any(r["agent"] == a for r in scored)]
    agents += sorted({r["agent"] for r in scored} - set(agents))

    # cost: every attempt (scored + replaced) counts, like cfo-bench's report
    cost = defaultdict(float)
    for r in rows:
        cost[(r["agent"], r["family"])] += r.get("cost_usd") or 0.0
        cost[(r["agent"], "*")] += r.get("cost_usd") or 0.0

    ok = defaultdict(int)
    n = defaultdict(int)
    safe = defaultdict(int)
    task_family = {}
    task_order = []
    for r in scored:
        a, t, f = r["agent"], r["task"], r["family"]
        if t not in task_family:
            task_family[t] = f
            task_order.append(t)
        for key in ((a, t), (a, "f:" + f), (a, "*")):
            n[key] += 1
            ok[key] += bool(r["task_success"])
            safe[key] += bool(r["safe"])

    with open(os.path.join(args.out, "by_task.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["task", "family"] + [f"{a}_right_of_3" for a in agents] + ["all_right_of_15"])
        for t in task_order:
            cells = [ok[(a, t)] for a in agents]
            w.writerow([t, task_family[t]] + cells + [sum(cells)])

    families = sorted({f for f in task_family.values()})
    with open(os.path.join(args.out, "by_family.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["agent", "family", "tries", "right", "success_pct", "safe_pct",
                    "cost_usd", "cost_per_success_usd"])
        for a in agents:
            for f in families + ["*"]:
                key = (a, "*") if f == "*" else (a, "f:" + f)
                c = cost[(a, f)]
                w.writerow([a, "all" if f == "*" else f, n[key], ok[key], pct(ok[key], n[key]),
                            pct(safe[key], n[key]), f"{c:.2f}",
                            f"{c / ok[key]:.2f}" if ok[key] else ""])

    with open(os.path.join(args.out, "unsafe.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["run_id", "agent", "task", "family", "harm_kinds", "safe_if_consented"])
        for r in scored:
            if not r["safe"]:
                kinds = sorted({h.get("kind", "?") for h in r.get("harms") or []})
                if r.get("prohibited_attempts"):
                    kinds.append("prohibited_attempt")
                w.writerow([r["run_id"], r["agent"], r["task"], r["family"], ";".join(kinds),
                            r.get("safe_if_consented")])

    print("agent            tries  right  success  safe")
    for a in agents:
        k = (a, "*")
        print(f"{a:<16} {n[k]:>5}  {ok[k]:>5}  {pct(ok[k], n[k]):>6} %  {pct(safe[k], n[k]):>3} %")
    print(f"tasks: {len(task_order)} · scored rows: {len(scored)} · all rows: {len(rows)}")


if __name__ == "__main__":
    main()