#!/usr/bin/env python3
"""cc-tokens — where your Claude Code quota actually went.

Reads the session transcripts Claude Code writes to ~/.claude/projects/**/*.jsonl
and attributes token spend by day, session, project, model, and token *category*.
Pure stdlib, fully offline, no network and no npm.

The headline it exists to surface: on a long-running agentic session, almost all
token spend is CONTEXT RE-READ (cache_read), not generated output. Cost scales
with (session length x context size), largely independent of work produced.

Subcommands
  daily        per-day totals + composition
  sessions     per-session, ranked — the culprits, with project and lifespan
  projects     per-project rollup
  models       per-model rollup
  composition  where the tokens/cost go by category (the "why")
  hogs         diagnosis: flags marathon sessions, bloated contexts, fan-out
  context      per-turn context-growth curve for ONE session

Common flags
  --since YYYY-MM-DD   inclusive lower bound (default: 7 days back)
  --until YYYY-MM-DD   inclusive upper bound
  --root PATH          projects dir (default ~/.claude/projects)
  --json               machine-readable output
  --top N              rows to show (default 15)

Cost is an API-equivalent estimate, NOT subscription quota. Use it to RANK
consumers, not to predict when you hit the weekly cap — for the real quota bars
run /usage inside Claude Code.
"""

import argparse
import json
import os
import sys
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from pathlib import Path

# ---------------------------------------------------------------- pricing ----
# USD per million tokens, as (input, output). Cache multipliers are applied on
# top of the INPUT rate and are uniform across Claude models:
#   cache read        0.1x input
#   cache write 5m    1.25x input
#   cache write 1h    2.0x input     <- Claude Code uses the 1h TTL
# Source: platform.claude.com pricing + prompt-caching docs (checked 2026-07-22).
CACHE_READ_MULT = 0.10
CACHE_WRITE_5M_MULT = 1.25
CACHE_WRITE_1H_MULT = 2.00

PRICES = {
    "claude-fable-5": (10.0, 50.0),
    "claude-mythos-5": (10.0, 50.0),
    "claude-mythos-preview": (10.0, 50.0),
    "claude-opus-4-8": (5.0, 25.0),
    "claude-opus-4-7": (5.0, 25.0),
    "claude-opus-4-6": (5.0, 25.0),
    "claude-opus-4-5": (5.0, 25.0),
    "claude-opus-4-1": (15.0, 75.0),
    "claude-opus-4-0": (15.0, 75.0),
    # Sonnet 5 list is 3/15; an introductory 2/10 applies through 2026-08-31.
    "claude-sonnet-5": (2.0, 10.0),
    "claude-sonnet-4-6": (3.0, 15.0),
    "claude-sonnet-4-5": (3.0, 15.0),
    "claude-sonnet-4-0": (3.0, 15.0),
    "claude-haiku-4-5": (1.0, 5.0),
    "claude-3-haiku": (0.25, 1.25),
}
# Tier weight for "opus-equivalent" normalisation (input $/Mtok / 5.0).
DEFAULT_PRICE = (5.0, 25.0)


def price_for(model):
    """Resolve a model id to (input, output) $/Mtok, tolerating date suffixes."""
    if not model:
        return DEFAULT_PRICE
    if model in PRICES:
        return PRICES[model]
    # Strip a trailing -YYYYMMDD snapshot suffix, then try longest prefix match.
    best = None
    for known in PRICES:
        if model.startswith(known) and (best is None or len(known) > len(best)):
            best = known
    return PRICES[best] if best else DEFAULT_PRICE


class Rec:
    """One billed assistant turn."""

    __slots__ = ("ts", "model", "session", "cwd", "sidechain", "inp", "cw5",
                 "cw1", "cr", "out")

    def __init__(self, ts, model, session, cwd, sidechain, inp, cw5, cw1, cr, out):
        self.ts, self.model, self.session, self.cwd = ts, model, session, cwd
        self.sidechain = sidechain
        self.inp, self.cw5, self.cw1, self.cr, self.out = inp, cw5, cw1, cr, out

    @property
    def total(self):
        return self.inp + self.cw5 + self.cw1 + self.cr + self.out

    @property
    def cost(self):
        pin, pout = price_for(self.model)
        return (
            self.inp * pin
            + self.cw5 * pin * CACHE_WRITE_5M_MULT
            + self.cw1 * pin * CACHE_WRITE_1H_MULT
            + self.cr * pin * CACHE_READ_MULT
            + self.out * pout
        ) / 1e6

    @property
    def weighted(self):
        """Input-token-equivalents, tier-normalised to opus (pricing-free-ish).

        Useful when you want a single comparable number across models without
        arguing about dollars. weighted = tier * (in + 1.25*cw5m + 2*cw1h
        + 0.1*cr + 5*out), tier = input_rate / 5.
        """
        pin, _ = price_for(self.model)
        return (pin / 5.0) * (
            self.inp
            + self.cw5 * CACHE_WRITE_5M_MULT
            + self.cw1 * CACHE_WRITE_1H_MULT
            + self.cr * CACHE_READ_MULT
            + self.out * 5.0
        )

    @property
    def ctx(self):
        """Context size actually sent on this turn (everything but output)."""
        return self.inp + self.cw5 + self.cw1 + self.cr


# ------------------------------------------------------------------ load ----

def parse_day(s):
    if not s:
        return None
    s = s.strip().replace("/", "-")
    for fmt in ("%Y-%m-%d", "%Y%m%d"):
        try:
            return datetime.strptime(s, fmt).replace(tzinfo=timezone.utc)
        except ValueError:
            pass
    raise SystemExit(f"cc-tokens: cannot parse date {s!r} (want YYYY-MM-DD)")


def load(root, since=None, until=None, session=None):
    """Yield one Rec per billed assistant turn.

    Two distinct de-duplications, and getting either wrong silently skews
    everything downstream:

    1. One message spans MANY jsonl lines. Claude Code appends a line per content
       block (thinking, text, each tool_use), and every line repeats the SAME
       message-level `usage`. Summing lines inflates a turn several-fold.
    2. `output_tokens` GROWS across those lines — they are streaming snapshots of
       one response. Only the final line is complete, so we take the max per key
       rather than the first. (Verified empirically: output is monotonically
       non-decreasing within a key; the cache fields never vary.)

    The key is (message.id, requestId), which is also stable across files, so a
    resumed / forked / worktree-replayed session does not double count.
    """
    root = Path(root).expanduser()
    if not root.is_dir():
        raise SystemExit(f"cc-tokens: no such directory: {root}")

    # mtime prefilter: a file last touched before `since` cannot contain rows in
    # range. Cheap, and on a 1.6 GB history it is the difference between 4s and 40s.
    floor = since.timestamp() - 86400 if since else None
    seen = set()

    for path in root.rglob("*.jsonl"):
        if session and not path.stem.startswith(session):
            continue
        try:
            if floor is not None and path.stat().st_mtime < floor:
                continue
        except OSError:
            continue
        try:
            fh = path.open("r", errors="replace")
        except OSError:
            continue
        # Collapse this file's streaming snapshots first, then emit. Duplicates
        # are almost always intra-file, so a per-file map keeps memory bounded
        # while `seen` still catches the cross-file (resume/fork) case.
        pending = {}
        with fh:
            for line in fh:
                # Fast reject before paying for json.loads on ~1M lines.
                if '"usage"' not in line:
                    continue
                try:
                    o = json.loads(line)
                except (ValueError, TypeError):
                    continue
                if o.get("type") != "assistant":
                    continue
                msg = o.get("message") or {}
                usage = msg.get("usage") or {}
                if not usage:
                    continue
                # "<synthetic>" is Claude Code's own placeholder turn (API errors,
                # interrupts). Zero tokens, never billed — but left in it makes a
                # session look like it switched models.
                if msg.get("model") == "<synthetic>":
                    continue

                ts = o.get("timestamp")
                try:
                    when = datetime.fromisoformat(str(ts).replace("Z", "+00:00"))
                except (ValueError, TypeError):
                    continue
                if since and when < since:
                    continue
                if until and when > until:
                    continue

                cw = usage.get("cache_creation") or {}
                cw5 = int(cw.get("ephemeral_5m_input_tokens") or 0)
                cw1 = int(cw.get("ephemeral_1h_input_tokens") or 0)
                total_cw = int(usage.get("cache_creation_input_tokens") or 0)
                if not (cw5 or cw1) and total_cw:
                    # Older transcripts lack the TTL split. Claude Code writes at
                    # 1h, so attribute there rather than silently under-costing.
                    cw1 = total_cw

                rec = Rec(
                    ts=when,
                    model=msg.get("model"),
                    session=o.get("sessionId") or path.stem,
                    cwd=o.get("cwd") or "",
                    sidechain=bool(o.get("isSidechain")),
                    inp=int(usage.get("input_tokens") or 0),
                    cw5=cw5,
                    cw1=cw1,
                    cr=int(usage.get("cache_read_input_tokens") or 0),
                    out=int(usage.get("output_tokens") or 0),
                )

                key = (msg.get("id"), o.get("requestId"))
                if key == (None, None):
                    # No identity to collapse on — treat as its own turn.
                    key = ("_anon", id(rec))
                prev = pending.get(key)
                if prev is None:
                    pending[key] = rec
                else:
                    # Same message, later snapshot: keep the largest reading of
                    # each field. Max (not last) so out-of-order lines are safe.
                    for f in ("inp", "cw5", "cw1", "cr", "out"):
                        setattr(prev, f, max(getattr(prev, f), getattr(rec, f)))

        for key, rec in pending.items():
            if key in seen:
                continue
            seen.add(key)
            yield rec


# --------------------------------------------------------------- helpers ----

class Bucket:
    def __init__(self):
        self.inp = self.cw5 = self.cw1 = self.cr = self.out = 0
        self.cost = 0.0
        self.weighted = 0.0
        self.turns = 0
        self.sidechain_turns = 0
        self.models = set()
        self.first = None
        self.last = None
        self.max_ctx = 0
        self.projects = set()

    def add(self, r):
        self.inp += r.inp
        self.cw5 += r.cw5
        self.cw1 += r.cw1
        self.cr += r.cr
        self.out += r.out
        self.cost += r.cost
        self.weighted += r.weighted
        self.turns += 1
        if r.sidechain:
            self.sidechain_turns += 1
        if r.model:
            self.models.add(r.model)
        if r.cwd:
            self.projects.add(r.cwd)
        self.max_ctx = max(self.max_ctx, r.ctx)
        if self.first is None or r.ts < self.first:
            self.first = r.ts
        if self.last is None or r.ts > self.last:
            self.last = r.ts

    @property
    def total(self):
        return self.inp + self.cw5 + self.cw1 + self.cr + self.out

    def as_dict(self):
        return {
            "input": self.inp, "cache_write_5m": self.cw5,
            "cache_write_1h": self.cw1, "cache_read": self.cr,
            "output": self.out, "total_tokens": self.total,
            "cost_usd": round(self.cost, 2),
            "weighted_tokens": int(self.weighted),
            "turns": self.turns, "sidechain_turns": self.sidechain_turns,
            "models": sorted(self.models),
            "first": self.first.isoformat() if self.first else None,
            "last": self.last.isoformat() if self.last else None,
            "max_context": self.max_ctx,
        }


def human(n):
    for unit, div in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("k", 1e3)):
        if abs(n) >= div:
            return f"{n / div:.1f}{unit}"
    return str(int(n))


def group(records, keyfn):
    out = defaultdict(Bucket)
    for r in records:
        out[keyfn(r)].add(r)
    return out


def project_of(cwd):
    """Short label for a working directory."""
    if not cwd:
        return "?"
    p = Path(cwd)
    parts = p.parts
    if ".claude" in parts and "worktrees" in parts:
        i = parts.index("worktrees")
        base = parts[i - 2] if i >= 2 else p.name
        return f"{base} (worktree:{parts[-1]})"
    return p.name or str(p)


def table(rows, headers, aligns=None):
    if not rows:
        print("  (nothing in range)")
        return
    cols = len(headers)
    aligns = aligns or [">"] * cols
    cells = [[str(c) for c in row] for row in rows]
    widths = [max(len(headers[i]), max(len(r[i]) for r in cells)) for i in range(cols)]
    print("  " + "  ".join(h.ljust(widths[i]) if aligns[i] == "<" else h.rjust(widths[i])
                           for i, h in enumerate(headers)))
    print("  " + "  ".join("-" * widths[i] for i in range(cols)))
    for r in cells:
        print("  " + "  ".join(r[i].ljust(widths[i]) if aligns[i] == "<" else r[i].rjust(widths[i])
                               for i in range(cols)))


def emit_json(obj):
    json.dump(obj, sys.stdout, indent=2, default=str)
    print()


# ------------------------------------------------------------ subcommands ----

def cmd_daily(recs, args):
    g = group(recs, lambda r: r.ts.astimezone().strftime("%Y-%m-%d"))
    if args.json:
        return emit_json({k: v.as_dict() for k, v in sorted(g.items())})
    rows = []
    for day in sorted(g):
        b = g[day]
        rows.append([day, f"{b.cost:,.0f}", human(b.out), human(b.cw5 + b.cw1),
                     human(b.cr), human(b.total), b.turns])
    tot = Bucket()
    for b in g.values():
        for f in ("inp", "cw5", "cw1", "cr", "out", "turns"):
            setattr(tot, f, getattr(tot, f) + getattr(b, f))
        tot.cost += b.cost
    rows.append(["TOTAL", f"{tot.cost:,.0f}", human(tot.out), human(tot.cw5 + tot.cw1),
                 human(tot.cr), human(tot.total), tot.turns])
    print("\nPer-day spend (cost is API-equivalent USD, not quota)\n")
    table(rows, ["day", "$", "output", "cacheWrite", "cacheRead", "total", "turns"],
          ["<", ">", ">", ">", ">", ">", ">"])


def cmd_sessions(recs, args):
    g = group(recs, lambda r: r.session)
    ranked = sorted(g.items(), key=lambda kv: -kv[1].cost)
    if args.json:
        return emit_json({k: v.as_dict() for k, v in ranked[: args.top]})
    grand = sum(b.cost for _, b in ranked)
    rows = []
    for sid, b in ranked[: args.top]:
        span = ""
        if b.first and b.last:
            days = (b.last - b.first).days
            span = f"{b.first.astimezone():%m-%d}>{b.last.astimezone():%m-%d} ({days}d)"
        proj = project_of(sorted(b.projects)[0]) if b.projects else "?"
        rows.append([sid[:8], f"{b.cost:,.0f}", b.turns, human(b.max_ctx),
                     b.sidechain_turns, span, proj[:34]])
    print(f"\nTop sessions by spend — {len(ranked)} sessions, ${grand:,.0f} total\n")
    table(rows, ["session", "$", "turns", "maxCtx", "subagt", "lifespan", "project"],
          ["<", ">", ">", ">", ">", "<", "<"])
    shown = sum(b.cost for _, b in ranked[: args.top])
    if grand:
        print(f"\n  top {min(args.top, len(ranked))} = {100 * shown / grand:.1f}% of spend")


def cmd_projects(recs, args):
    g = group(recs, lambda r: project_of(r.cwd))
    ranked = sorted(g.items(), key=lambda kv: -kv[1].cost)
    if args.json:
        return emit_json({k: v.as_dict() for k, v in ranked})
    rows = [[name[:40], f"{b.cost:,.0f}", b.turns, human(b.cr), human(b.out)]
            for name, b in ranked[: args.top]]
    print("\nSpend by project\n")
    table(rows, ["project", "$", "turns", "cacheRead", "output"], ["<", ">", ">", ">", ">"])


def cmd_models(recs, args):
    g = group(recs, lambda r: r.model or "?")
    ranked = sorted(g.items(), key=lambda kv: -kv[1].cost)
    if args.json:
        return emit_json({k: v.as_dict() for k, v in ranked})
    grand = sum(b.cost for _, b in ranked) or 1.0
    rows = [[m, f"{b.cost:,.0f}", f"{100 * b.cost / grand:.1f}%", b.turns,
             human(b.total)] for m, b in ranked]
    print("\nSpend by model\n")
    table(rows, ["model", "$", "share", "turns", "tokens"], ["<", ">", ">", ">", ">"])


def cmd_composition(recs, args):
    b = Bucket()
    for r in recs:
        b.add(r)
    if not b.turns:
        print("  (nothing in range)")
        return
    pin, pout = 5.0, 25.0  # display-only reference tier for the cost column
    cats = [
        ("cache_read", b.cr, CACHE_READ_MULT * pin),
        ("cache_write_1h", b.cw1, CACHE_WRITE_1H_MULT * pin),
        ("cache_write_5m", b.cw5, CACHE_WRITE_5M_MULT * pin),
        ("output", b.out, pout),
        ("input", b.inp, pin),
    ]
    if args.json:
        d = b.as_dict()
        d["cost_by_category_usd"] = {n: round(t * rate / 1e6, 2) for n, t, rate in cats}
        return emit_json(d)
    costs = {n: t * rate / 1e6 for n, t, rate in cats}
    ctot = sum(costs.values()) or 1.0
    rows = []
    for name, toks, _ in cats:
        rows.append([name, human(toks), f"{100 * toks / (b.total or 1):.1f}%",
                     f"{100 * costs[name] / ctot:.1f}%"])
    print(f"\nToken composition — {b.turns:,} turns, ${b.cost:,.0f}\n")
    table(rows, ["category", "tokens", "%tokens", "%cost"], ["<", ">", ">", ">"])
    print(f"\n  context re-read (cache_read) is {100 * b.cr / (b.total or 1):.1f}% of all tokens")
    print(f"  generated output is only        {100 * b.out / (b.total or 1):.1f}%")
    mean_ctx = (b.total - b.out) / b.turns
    print(f"  mean context sent per turn:     {human(mean_ctx)}")


def cmd_context(recs, args):
    if not args.session:
        raise SystemExit("cc-tokens context: --session <id-or-prefix> is required")
    rows = sorted((r for r in recs if r.session.startswith(args.session)),
                  key=lambda r: r.ts)
    if not rows:
        raise SystemExit(f"cc-tokens: no turns for session {args.session!r} in range")
    if args.json:
        return emit_json([{"ts": r.ts.isoformat(), "context": r.ctx,
                           "cache_read": r.cr, "output": r.out,
                           "cost_usd": round(r.cost, 4)} for r in rows])
    print(f"\nContext growth — session {rows[0].session[:8]}, {len(rows):,} turns\n")
    # Bucket into 20 slices so the shape is visible regardless of turn count.
    n = min(20, len(rows))
    size = max(1, len(rows) // n)
    peak = max(r.ctx for r in rows) or 1
    out = []
    for i in range(0, len(rows), size):
        chunk = rows[i:i + size]
        mean = sum(c.ctx for c in chunk) / len(chunk)
        bar = "#" * int(40 * mean / peak)
        out.append([f"{i + 1}-{i + len(chunk)}", human(mean),
                    f"{sum(c.cost for c in chunk):,.2f}", bar])
    table(out, ["turns", "avgCtx", "$", ""], ["<", ">", ">", "<"])
    tot = sum(r.cost for r in rows)
    print(f"\n  peak context {human(peak)} | total ${tot:,.0f} | "
          f"${tot / len(rows):.2f}/turn")


def cmd_hogs(recs, args):
    """The opinionated one: name the specific things worth changing."""
    recs = list(recs)
    if not recs:
        print("  (nothing in range)")
        return
    sess = group(recs, lambda r: r.session)
    grand = sum(b.cost for b in sess.values()) or 1.0
    flagged = []

    for sid, b in sess.items():
        share = 100 * b.cost / grand
        proj = project_of(sorted(b.projects)[0]) if b.projects else "?"
        days = (b.last - b.first).days if (b.first and b.last) else 0
        notes = []

        if days >= 2 and b.cost > grand * 0.02:
            notes.append(("marathon-session",
                          f"alive {days}d across {b.turns:,} turns — every turn re-reads "
                          f"the whole context; /clear and start fresh"))
        if b.max_ctx > 300_000:
            notes.append(("bloated-context",
                          f"peak context {human(b.max_ctx)}, re-sent on every turn — "
                          f"trim MCP servers / CLAUDE.md / pasted files"))
        if b.turns and b.sidechain_turns / b.turns > 0.5 and b.cost > grand * 0.02:
            notes.append(("subagent-heavy",
                          f"{100 * b.sidechain_turns / b.turns:.0f}% of turns are subagents "
                          f"— each carries its own full context"))
        if len(b.models) > 1 and b.cost > grand * 0.05:
            short = ", ".join(sorted(m.replace("claude-", "") for m in b.models))
            notes.append(("model-switching",
                          f"models {short} in one session — a switch invalidates "
                          f"the whole cached prefix"))
        if notes:
            flagged.append((b.cost, share, sid[:8], proj, b, notes))

    flagged.sort(key=lambda f: -f[0])
    if args.json:
        return emit_json([
            {"session": s, "project": p, "cost_usd": round(c, 2),
             "share_pct": round(sh, 1), "turns": b.turns,
             "max_context": b.max_ctx,
             "findings": [{"kind": k, "detail": d} for k, d in notes]}
            for c, sh, s, p, b, notes in flagged])

    tot = Bucket()
    for r in recs:
        tot.add(r)
    print(f"\nDiagnosis — {len(sess)} sessions, ${grand:,.0f}, "
          f"{100 * tot.cr / (tot.total or 1):.0f}% of all tokens are context re-read\n")
    if not flagged:
        print("  No structural problems found. Spend looks proportional to work.")
        return
    covered = 0.0
    for cost, share, sid, proj, b, notes in flagged[: args.top]:
        covered += cost
        print(f"  {sid}  {proj[:32]}   ${cost:,.0f}  ({share:.1f}% of spend, "
              f"{b.turns:,} turns)")
        for kind, detail in notes:
            print(f"      - {kind}: {detail}")
        print()
    print(f"  {min(args.top, len(flagged))} flagged sessions = "
          f"{100 * covered / grand:.0f}% of spend. Fix from the top.")


COMMANDS = {
    "daily": cmd_daily, "sessions": cmd_sessions, "projects": cmd_projects,
    "models": cmd_models, "composition": cmd_composition, "hogs": cmd_hogs,
    "context": cmd_context,
}


def main(argv=None):
    ap = argparse.ArgumentParser(
        prog="cc-tokens", description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("command", choices=sorted(COMMANDS), nargs="?", default="hogs")
    ap.add_argument("--since", help="inclusive start date YYYY-MM-DD (default: 7d ago)")
    ap.add_argument("--until", help="inclusive end date YYYY-MM-DD")
    ap.add_argument("--days", type=int, help="shorthand for --since N days ago")
    ap.add_argument("--session", help="session id or prefix (required for `context`)")
    ap.add_argument("--root", default=os.environ.get(
        "CC_TOKENS_ROOT", str(Path.home() / ".claude" / "projects")))
    ap.add_argument("--top", type=int, default=15)
    ap.add_argument("--json", action="store_true")
    args = ap.parse_args(argv)

    if args.days:
        since = datetime.now(timezone.utc) - timedelta(days=args.days)
    elif args.since:
        since = parse_day(args.since)
    else:
        since = datetime.now(timezone.utc) - timedelta(days=7)
    until = parse_day(args.until) + timedelta(days=1) if args.until else None

    recs = load(args.root, since=since, until=until,
                session=args.session if args.command == "context" else None)
    COMMANDS[args.command](recs, args)


if __name__ == "__main__":
    try:
        main()
    except BrokenPipeError:
        sys.stderr.close()
    except KeyboardInterrupt:
        sys.exit(130)
