#!/usr/bin/env python3
"""Runway Audit: where do your Claude Code sessions spend tokens?

Reads the local transcript files Claude Code already writes
(~/.claude/projects/**/*.jsonl). It reads token counts, model names,
timestamps and compaction markers only. It never reads what you or
Claude wrote, never uses the network, and needs no API key.

Usage:
  python3 runway_audit.py                     # last 28 days
  python3 runway_audit.py --days 14
  python3 runway_audit.py --since 2026-09-01 --until 2026-09-29
  python3 runway_audit.py --project .         # also size CLAUDE.md files here
  python3 runway_audit.py --json > baseline.json

Requires Python 3.8 or newer. Standard library only.
"""
import argparse
import glob
import json
import os
import statistics
import sys
from datetime import datetime, timedelta, timezone

VERSION = "1.0.0"
CLASSES = ("input", "cache_write", "cache_read", "output")
LONG_SESSION_TURNS = 40

# Fit rule: an uncalibrated screening heuristic, not a prediction of benefit.
# The thresholds were chosen by someone who already knew their own history was
# heavy-tailed; see the method section on chudi.dev/tools/claude-code-runway-audit.
MIN_SESSIONS = 10
MIN_TURNS = 300
FIT_LONG_SESSION_SHARE = 0.50
FIT_COMPACTIONS_PER_10 = 1.0


def parse_ts(s):
    try:
        return datetime.fromisoformat(s.replace("Z", "+00:00"))
    except (AttributeError, ValueError):
        return None


def parse_day(s):
    return datetime.strptime(s, "%Y-%m-%d").replace(tzinfo=timezone.utc)


def read_sessions(root, since, until):
    """Return {session_key: session dict} from transcript files under root."""
    sessions = {}
    seen_ids = set()
    versions = set()
    oldest = None
    files = glob.glob(os.path.join(root, "**", "*.jsonl"), recursive=True)
    for path in files:
        try:
            if datetime.fromtimestamp(os.path.getmtime(path), timezone.utc) < since:
                continue  # file untouched since the window opened
        except OSError:
            continue
        is_sub = os.sep + "subagents" + os.sep in path
        key = path
        s = None
        try:
            fh = open(path, encoding="utf-8", errors="replace")
        except OSError:
            continue
        with fh:
            for line in fh:
                # Cheap prefilter: only usage rows and compaction markers matter.
                if '"usage"' not in line and "compact_boundary" not in line:
                    continue
                try:
                    row = json.loads(line)
                except ValueError:
                    continue
                if not isinstance(row, dict):
                    continue
                if row.get("version"):
                    versions.add(row["version"])
                ts = parse_ts(row.get("timestamp"))
                if ts is not None and (oldest is None or ts < oldest):
                    oldest = ts
                if ts is None or ts < since or ts >= until:
                    continue
                if s is None:
                    s = sessions.setdefault(key, {"subagent": is_sub, "turns": [], "compactions": []})
                if row.get("type") == "system" and row.get("subtype") == "compact_boundary":
                    meta = row.get("compactMetadata") or {}
                    s["compactions"].append(meta.get("preTokens"))
                    continue
                msg = row.get("message")
                if row.get("type") != "assistant" or not isinstance(msg, dict):
                    continue
                usage = msg.get("usage")
                if not isinstance(usage, dict):
                    continue
                # One API response can be split across several rows that repeat
                # the same usage; count each message id once.
                mid = msg.get("id") or row.get("requestId")
                if mid:
                    if mid in seen_ids:
                        continue
                    seen_ids.add(mid)
                s["turns"].append({
                    "model": msg.get("model") or "unknown",
                    "input": int(usage.get("input_tokens") or 0),
                    "cache_write": int(usage.get("cache_creation_input_tokens") or 0),
                    "cache_read": int(usage.get("cache_read_input_tokens") or 0),
                    "output": int(usage.get("output_tokens") or 0),
                })
    return sessions, len(files), sorted(versions), oldest


def total(turn):
    return sum(turn[c] for c in CLASSES)


def share(part, whole):
    return round(part / whole, 4) if whole else None


def analyze(sessions):
    main = [s for s in sessions.values() if s["turns"] and not s["subagent"]]
    subs = [s for s in sessions.values() if s["turns"] and s["subagent"]]
    all_turns = [t for s in main + subs for t in s["turns"]]
    grand = sum(total(t) for t in all_turns)
    by_class = {c: sum(t[c] for t in all_turns) for c in CLASSES}
    by_model = {}
    for t in all_turns:
        by_model[t["model"]] = by_model.get(t["model"], 0) + total(t)

    # Concentration over main sessions (subagent tokens counted separately).
    sizes = sorted((sum(total(t) for t in s["turns"]) for s in main), reverse=True)
    top_n = max(1, len(sizes) // 10) if sizes else 0
    top10 = share(sum(sizes[:top_n]), sum(sizes)) if sizes else None

    long_tokens = sum(sum(total(t) for t in s["turns"]) for s in main if len(s["turns"]) >= LONG_SESSION_TURNS)
    turn_counts = [len(s["turns"]) for s in main]
    ratios = []
    for s in main:
        n = len(s["turns"])
        if n < LONG_SESSION_TURNS:
            continue
        q = n // 4
        first = statistics.median(t["cache_read"] for t in s["turns"][:q])
        last = statistics.median(t["cache_read"] for t in s["turns"][-q:])
        if first > 0:
            ratios.append(last / first)
    compactions = [c for s in main for c in s["compactions"]]
    pre = [c for c in compactions if isinstance(c, (int, float))]

    return {
        "sessions": len(main),
        "subagent_runs": len(subs),
        "assistant_turns": len(all_turns),
        "tokens_total": grand,
        "tokens_by_class": by_class,
        "share_by_class": {c: share(v, grand) for c, v in by_class.items()},
        "subagent_token_share": share(sum(total(t) for s in subs for t in s["turns"]), grand),
        "share_by_model": {m: share(v, grand) for m, v in sorted(by_model.items(), key=lambda kv: -kv[1])},
        "top10pct_session_share": top10,
        "top10pct_session_count": top_n,
        "turns_per_session_median": statistics.median(turn_counts) if turn_counts else None,
        "turns_per_session_p90": sorted(turn_counts)[int(0.9 * (len(turn_counts) - 1))] if turn_counts else None,
        "long_session_token_share": share(long_tokens, sum(sizes)) if sizes else None,
        "long_sessions": len(ratios),
        "late_vs_early_cache_read_per_turn_median": round(statistics.median(ratios), 2) if ratios else None,
        "compactions": len(compactions),
        "compactions_per_10_sessions": round(10 * len(compactions) / len(main), 2) if main else None,
        "compaction_pre_tokens_median": int(statistics.median(pre)) if pre else None,
    }


def fit(r):
    if r["sessions"] < MIN_SESSIONS or r["assistant_turns"] < MIN_TURNS:
        return "INSUFFICIENT DATA", [f"needs at least {MIN_SESSIONS} sessions and {MIN_TURNS} turns in the window"]
    reasons = []
    if (r["long_session_token_share"] or 0) >= FIT_LONG_SESSION_SHARE:
        reasons.append(f"sessions of {LONG_SESSION_TURNS}+ turns carry most of your tokens")
    if (r["compactions_per_10_sessions"] or 0) >= FIT_COMPACTIONS_PER_10:
        reasons.append("your sessions hit compaction regularly")
    return {2: "LIKELY FIT", 1: "POSSIBLE FIT", 0: "PROBABLY NOT A FIT"}[len(reasons)], reasons


def claude_md_sizes(project, config_dir):
    paths = [os.path.join(config_dir, "CLAUDE.md")]
    if project:
        paths += glob.glob(os.path.join(project, "**", "CLAUDE.md"), recursive=True)
        paths += glob.glob(os.path.join(project, "CLAUDE.local.md"))
    out = []
    for p in sorted(set(paths)):
        try:
            chars = len(open(p, encoding="utf-8", errors="replace").read())
        except OSError:
            continue
        out.append({"path": p, "chars": chars, "approx_tokens": chars // 4})
    return out


def pct(x):
    return "n/a" if x is None else f"{100 * x:.1f}%"


def main(argv=None):
    ap = argparse.ArgumentParser(description="Runway Audit for Claude Code transcripts (local, read-only).")
    config_dir = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
    ap.add_argument("--root", default=os.path.join(config_dir, "projects"))
    ap.add_argument("--days", type=int, default=28)
    ap.add_argument("--since", help="YYYY-MM-DD (UTC, inclusive)")
    ap.add_argument("--until", help="YYYY-MM-DD (UTC, exclusive)")
    ap.add_argument("--project", help="also size CLAUDE.md files under this folder")
    ap.add_argument("--json", action="store_true", help="print JSON (save it as your baseline)")
    a = ap.parse_args(argv)

    until = parse_day(a.until) if a.until else datetime.now(timezone.utc)
    since = parse_day(a.since) if a.since else until - timedelta(days=a.days)
    sessions, nfiles, versions, oldest = read_sessions(a.root, since, until)
    r = analyze(sessions)
    verdict, reasons = fit(r)
    report = {
        "tool": "runway_audit", "tool_version": VERSION,
        "window": {"since": since.isoformat(), "until": until.isoformat()},
        "root": a.root, "files_scanned": nfiles, "claude_code_versions_seen": versions[-5:],
        "oldest_row_in_changed_files": oldest.isoformat() if oldest else None,
        "status": "OK" if r["assistant_turns"] else "UNREADABLE",
        "results": r, "fit": {"verdict": verdict, "reasons": reasons},
        "claude_md": claude_md_sizes(a.project, config_dir),
    }
    if a.json:
        print(json.dumps(report, indent=2))
        return 0 if report["status"] == "OK" else 2

    print(f"Runway Audit {VERSION}  window {since:%Y-%m-%d} to {until:%Y-%m-%d} (UTC)")
    if report["status"] == "UNREADABLE":
        print(f"\nUNREADABLE: scanned {nfiles} transcript files and found no token usage in this window.")
        print("This is not zero usage. Either the window is empty or your Claude Code version")
        print(f"writes a format this version of the audit does not know (seen: {', '.join(versions[-3:]) or 'none'}).")
        return 2
    print(f"\nCounted: {r['sessions']} sessions, {r['subagent_runs']} subagent runs, {r['assistant_turns']} assistant turns")
    print(f"Tokens: {r['tokens_total']:,}")
    print("\nWhere the tokens went")
    for c in CLASSES:
        print(f"  {c:<12} {r['tokens_by_class'][c]:>16,}  {pct(r['share_by_class'][c])}")
    print(f"  subagents carried {pct(r['subagent_token_share'])} of all tokens")
    print("\nBy model")
    for m, v in list(r["share_by_model"].items())[:6]:
        print(f"  {m:<32} {pct(v)}")
    print("\nSession shape")
    print(f"  top 10% of sessions ({r['top10pct_session_count']}) carry {pct(r['top10pct_session_share'])} of session tokens")
    print(f"  sessions of {LONG_SESSION_TURNS}+ turns carry {pct(r['long_session_token_share'])} of session tokens")
    print(f"  turns per session: median {r['turns_per_session_median']}, 90th percentile {r['turns_per_session_p90']}")
    if r["long_sessions"]:
        print(f"  in {r['long_sessions']} sessions of {LONG_SESSION_TURNS}+ turns, a late turn re-reads "
              f"{r['late_vs_early_cache_read_per_turn_median']}x the cached context of an early turn (median;")
        print("  this grows in any long session by design, so it shows the cost of staying, not a fault)")
    print(f"  compactions: {r['compactions']} ({r['compactions_per_10_sessions']} per 10 sessions)"
          + (f", median size before compaction {r['compaction_pre_tokens_median']:,} tokens" if r["compaction_pre_tokens_median"] else ""))
    if report["claude_md"]:
        print("\nLoaded every turn (CLAUDE.md, approximate tokens = characters / 4)")
        for f in report["claude_md"]:
            print(f"  {f['approx_tokens']:>7,}  {f['path']}")
    print(f"\nFit: {verdict}  (a rough screening rule, not a prediction of savings)")
    for why in reasons:
        print(f"  - {why}")
    if verdict == "PROBABLY NOT A FIT":
        print("  Your sessions are already short and light. The Runway Kit would probably not change much; save your money.")
    elif verdict in ("LIKELY FIT", "POSSIBLE FIT"):
        print("  This audit is free and stays free. The paid kit adds the change steps and the")
        print("  before/after comparison, not more measurement.")
    if report["oldest_row_in_changed_files"]:
        print(f"\nOldest transcript row found: {report['oldest_row_in_changed_files'][:10]}")
    print("Save this now: python3 runway_audit.py --json > runway-baseline.json")
    print("Claude Code deletes old transcripts, so this file may be the only record of your before period.")
    print("\nLimits: token counts are what Claude Code logged on this machine. Plan limits weight tokens")
    print("in ways Anthropic does not publish, so this shows shape, not your remaining allowance.")
    return 0


if __name__ == "__main__":
    sys.exit(main())