"""Reanalyse frozen historical tau-bench records; stdlib only. No model calls.
Run: python3 analyse.py
Counts assistant messages as interactions, NOT tokens, latency or dollars.
"""
import collections, hashlib, json, urllib.request, math
from pathlib import Path
ROOT = Path(__file__).resolve().parent
SHA = "59a200c6d575d595120f1cb70fea53cef0632f6b"
BASE = f"https://raw.githubusercontent.com/sierra-research/tau-bench/{SHA}/"
EXPECTED_HASHES = {'gpt-4o-airline.json': 'e9e6c0297660c537f83d4fd9c476ce7a9a86ecd2784874b7bfc13be598e37bfa', 'gpt-4o-retail.json': 'df01707894836168ff0ec9616b0bf08f66c7e5afcf313e5fe4f7a2f5c2ec938b', 'sonnet-35-new-airline.json': 'fe62fcd514b855b36f156dd4c3c7748597b392b006aff739b53337a9f3ba94d1', 'sonnet-35-new-retail.json': '0df526398e9d2720c32d340815cffb04fe8c4f8a61b1f4f84bf3bb558f760131'}
FILES = ["gpt-4o-airline", "gpt-4o-retail", "sonnet-35-new-airline", "sonnet-35-new-retail"]
raw = ROOT / "data" / "raw"
raw.mkdir(parents=True, exist_ok=True)
results, manifest = [], []
for name in FILES:
    path = raw / (name + ".json")
    url = BASE + "historical_trajectories/" + path.name
    if not path.exists():
        path.write_bytes(urllib.request.urlopen(url).read())
    payload = path.read_bytes()
    if hashlib.sha256(payload).hexdigest() != EXPECTED_HASHES[path.name]:
        raise ValueError(f"Source hash mismatch: {path.name}")
    rows = json.loads(payload)
    seen, groups = set(), collections.defaultdict(list)
    interactions = failed_interactions = successes = 0
    for row in rows:
        key = (row["task_id"], row["trial"])
        assert key not in seen, f"Duplicate task/trial: {key}"
        seen.add(key)
        assert row["reward"] in (0, 1), "Unexpected nonbinary reward"
        turns = sum(m.get("role") == "assistant" for m in row["traj"])
        interactions += turns
        successes += row["reward"] == 1
        if row["reward"] == 0:
            failed_interactions += turns
        groups[row["task_id"]].append(row["reward"])
    trial_counts = sorted(set(map(len, groups.values())))
    assert len(trial_counts) == 1 and trial_counts[0] >= 4, trial_counts
    results.append(dict(dataset=name, attempts=len(rows), tasks=len(groups),
        trials_per_task=trial_counts[0], successful_attempts=successes,
        success_rate=successes/len(rows),
        pass4=sum(math.comb(int(sum(v)),4)/math.comb(len(v),4) for v in groups.values())/len(groups),
        all_observed_success=sum(all(v) for v in groups.values()),
        at_least_one_success=sum(any(v) for v in groups.values()),
        assistant_messages=interactions, failed_assistant_messages=failed_interactions,
        messages_per_attempt=interactions/len(rows),
        messages_per_success=interactions/successes if successes else None,
        failed_message_share=failed_interactions/interactions,
        agent_total_cost_present=sum("total_cost" in r for r in rows)))
    manifest.append(dict(url=url, sha256=hashlib.sha256(payload).hexdigest(), bytes=len(payload)))
license_path = raw / "SIERRA-LICENSE.txt"
if not license_path.exists():
    license_path.write_bytes(urllib.request.urlopen(BASE + "LICENSE").read())
out = dict(source_commit=SHA, kind="historical trace reanalysis", datasets=results, files=manifest)
(ROOT / "data" / "analysis.json").write_text(json.dumps(out, indent=2) + "\n")
print(json.dumps(results, indent=2))