#!/usr/bin/env bash
# Measure one run: total tokens (sum of per-message usage in the pi session file),
# wall clock (from /usr/bin/time -v output), and calls to the refactoring CLI.
# usage: bash measure_run.sh <run_dir>
set -uo pipefail
RUN_DIR="$1"

SESSION=$(find "$RUN_DIR/session" -name '*.jsonl' 2>/dev/null | head -1)
echo "--- Run $(basename "$RUN_DIR") ---"

if [ -z "$SESSION" ]; then
  echo "session file: MISSING"
else
  echo "session file: $SESSION"
  python3 - "$SESSION" <<'PY'
import json, sys

path = sys.argv[1]
usage_keys = ("input", "output", "cacheRead", "cacheWrite", "reasoning")
totals = dict.fromkeys(usage_keys, 0)
n_assistant = 0
n_tool_uses = 0
tool_counts = {}
first_ts = last_ts = None
model = None

def walk(obj):
    """Yield every dict in a JSON structure (sessions nest messages in lists)."""
    if isinstance(obj, dict):
        yield obj
        for v in obj.values():
            yield from walk(v)
    elif isinstance(obj, list):
        for v in obj:
            yield from walk(v)

for line in open(path):
    try:
        rec = json.loads(line)
    except json.JSONDecodeError:
        continue
    if first_ts is None:
        first_ts = rec.get("timestamp")
    last_ts = rec.get("timestamp") or last_ts
    if rec.get("type") == "model_change":
        model = rec.get("modelId")
    for d in walk(rec.get("message") or {}):
        role = d.get("role")
        if role != "assistant":
            continue
        if d.get("usage"):
            n_assistant += 1
            u = d["usage"]
            for k in usage_keys:
                totals[k] += u.get(k) or 0
        for c in d.get("content") or []:
            if isinstance(c, dict) and c.get("type") == "toolCall":
                n_tool_uses += 1
                name = c.get("name") or "?"
                tool_counts[name] = tool_counts.get(name, 0) + 1

print(f"model:                {model}")
if first_ts and last_ts:
    from datetime import datetime
    fmt = "%Y-%m-%dT%H:%M:%S.%f%z"
    t0 = datetime.strptime(first_ts, fmt)
    t1 = datetime.strptime(last_ts, fmt)
    print(f"session span:         {(t1 - t0).total_seconds():.1f} s")
print(f"assistant messages:   {n_assistant}")
print(f"tool calls:           {n_tool_uses} total")
for name, n in sorted(tool_counts.items(), key=lambda kv: -kv[1]):
    print(f"  {name:24} {n}")
print(f"tokens input:         {totals['input']}")
print(f"tokens output:        {totals['output']}")
print(f"tokens cacheRead:     {totals['cacheRead']}")
print(f"tokens cacheWrite:    {totals['cacheWrite']}")
print(f"tokens reasoning:     {totals['reasoning']}")
print(f"tokens TOTAL (sum):   {sum(totals.values())}")
PY
fi

if [ -f "$RUN_DIR/time.txt" ]; then
  echo "--- /usr/bin/time -v (wall clock + max RSS) ---"
  grep -E "Elapsed \(wall clock\)|Maximum resident set size|Exit status" "$RUN_DIR/time.txt"
fi

# crude proxy: how often the CLI was invoked during the run (from bash history absent;
# the session transcript already counted tool calls, so this is just informational)
if [ -f "$RUN_DIR/pi-stdout.txt" ]; then
  echo "--- pi stdout (last 15 lines) ---"
  tail -15 "$RUN_DIR/pi-stdout.txt"
fi