#!/usr/bin/env python3
"""Claude Usage Audit - local, privacy-first transcript scanner.

Runs on the CLIENT's machine. Reads Claude Code transcripts (~/.claude/projects/**/*.jsonl)
and writes ONLY aggregate numbers: no prompt text, no file paths, no project names, no code.
The client reviews summary.json and sends it to us; raw transcripts never leave their machine.

Usage:
    python claude_usage_audit.py [--root ~/.claude/projects] [--out ./audit-out] [--since YYYY-MM-DD]

Outputs (in --out):
    summary.json   aggregate stats (this is the only file the client sends)
    summary.md     human-readable version of the same numbers

Cost is reported in "input-equivalent tokens" (relative weights below), not dollars, so the
numbers stay valid when list prices change. Convert to dollars at delivery time.
"""
import argparse
import collections
import hashlib
import json
import os
import re
import sys

# Relative price weights vs. 1 uncached input token (public ratios; adjust at delivery time).
W_IN, W_CACHE_READ, W_CACHE_WRITE, W_OUT = 1.0, 0.1, 1.25, 5.0

# Short "nudge" prompts: the user is restarting work the agent paused. Allowlist only -
# anything else short is counted in an anonymous bucket and never echoed.
NUDGE_RE = re.compile(
    r"^(go|go on|continue|keep going|proceed|do it|do it yourself|just do it|try again|retry|"
    r"deploy|ship it|yes|yep|ok|okay|sure|y|1|2|and\??|more|next|fix it|finish|"
    r"dale|sigue|continua|continúa|hazlo|hazlo tu|hazlo tú|intenta de nuevo|reintenta|si|sí)[.!? ]*$",
    re.I,
)
CMD_RE = re.compile(r"<command-name>/?([^<]+)</command-name>")
# Slash commands are product names, not secrets; custom ones are hashed below.
BUILTIN_SLASH = {"clear", "model", "compact", "init", "review", "help", "resume", "cost",
                 "context", "config", "permissions", "agents", "mcp", "memory", "loop",
                 "schedule", "fast", "usage", "login", "logout", "doctor", "status"}


def text_of(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "\n".join(c.get("text", "") for c in content
                         if isinstance(c, dict) and c.get("type") == "text")
    return ""


def anon(name, prefix):
    """Stable anonymous label for names that could identify a client system."""
    return f"{prefix}-{hashlib.sha1(name.encode()).hexdigest()[:6]}"


def scan(root, since):
    C = collections.Counter
    tools, skills, agent_types, agent_models, mcp = C(), C(), C(), C(), C()
    tool_err, slash = C(), C()
    msgs_by_model = C()
    tok = collections.defaultdict(C)
    sessions_by_day = C()
    nudges, short_other = C(), 0
    prompt_lens = []
    sessions = []
    interrupts = rejects = compactions = n_files = 0
    seen = set()

    for d, _, fs in os.walk(root):
        for f in fs:
            if not f.endswith(".jsonl"):
                continue
            path = os.path.join(d, f)
            norm = path.replace("\\", "/")
            is_sub = "/subagents/" in norm or f.startswith("agent-")
            id2name = {}
            first_ts, n_prompts, n_asst, ctx_peak = None, 0, 0, 0
            try:
                fh = open(path, "r", encoding="utf-8", errors="replace")
            except OSError:
                continue
            n_files += 1
            with fh:
                for line in fh:
                    try:
                        o = json.loads(line)
                    except ValueError:
                        continue
                    ts = o.get("timestamp")
                    if ts and not first_ts:
                        first_ts = ts
                    if since and ts and ts[:10] < since:
                        continue
                    typ = o.get("type")
                    if typ == "summary" or o.get("isCompactSummary") or o.get("subtype") == "compact_boundary":
                        compactions += 1
                        continue
                    msg = o.get("message") or {}
                    if typ == "assistant":
                        model = msg.get("model", "?")
                        content = msg.get("content") or []
                        for c in content if isinstance(content, list) else []:
                            if not (isinstance(c, dict) and c.get("type") == "tool_use"):
                                continue
                            name = c.get("name", "?")
                            id2name[c.get("id")] = name
                            inp = c.get("input") or {}
                            if name.startswith("mcp__"):
                                mcp[anon(name.split("__")[1], "mcp")] += 1
                                tools["mcp (any)"] += 1
                            else:
                                tools[("[sub] " if is_sub else "") + name] += 1
                            if name == "Skill":
                                skills[anon(str(inp.get("skill") or "?"), "skill")] += 1
                            if name in ("Agent", "Task"):
                                agent_types[anon(str(inp.get("subagent_type") or "general"), "agent")] += 1
                                agent_models[str(inp.get("model") or "(inherit)")] += 1
                        key = (msg.get("id"), o.get("requestId"))
                        u = msg.get("usage")
                        if u and key not in seen:
                            seen.add(key)
                            msgs_by_model[model] += 1
                            t = tok[model]
                            t["in"] += u.get("input_tokens") or 0
                            t["out"] += u.get("output_tokens") or 0
                            cr = u.get("cache_read_input_tokens") or 0
                            t["cache_read"] += cr
                            t["cache_write"] += u.get("cache_creation_input_tokens") or 0
                            if not is_sub:
                                n_asst += 1
                                ctx_peak = max(ctx_peak, cr + (u.get("input_tokens") or 0)
                                               + (u.get("cache_creation_input_tokens") or 0))
                    elif typ == "user":
                        content = msg.get("content")
                        for c in content if isinstance(content, list) else []:
                            if isinstance(c, dict) and c.get("type") == "tool_result" and c.get("is_error"):
                                r = c.get("content")
                                rt = r if isinstance(r, str) else text_of(r)
                                tool_err[id2name.get(c.get("tool_use_id"), "?")] += 1
                                if "doesn't want to proceed" in rt or "was rejected" in rt:
                                    rejects += 1
                        if is_sub or o.get("isMeta"):
                            continue
                        t = text_of(content).strip()
                        if not t:
                            continue
                        if "[Request interrupted by user" in t:
                            interrupts += 1
                            continue
                        for m in CMD_RE.findall(t):
                            m = m.strip()
                            slash[m if m in BUILTIN_SLASH else anon(m, "custom-cmd")] += 1
                        if t.startswith(("<local-command", "<command-message", "<system-reminder",
                                         "<task-notification", "Caveat:")):
                            continue
                        n_prompts += 1
                        prompt_lens.append(len(t))
                        if len(t) <= 30:
                            if NUDGE_RE.match(t):
                                nudges[t.lower().strip(".!? ")] += 1
                            else:
                                short_other += 1
            if is_sub or not first_ts:
                continue
            if since and first_ts[:10] < since:
                continue
            sessions_by_day[first_ts[:10]] += 1
            sessions.append({"prompts": n_prompts, "assistant_msgs": n_asst, "ctx_peak": ctx_peak})

    return dict(tools=tools, skills=skills, agent_types=agent_types, agent_models=agent_models,
                mcp=mcp, tool_err=tool_err, slash=slash, msgs_by_model=msgs_by_model, tok=tok,
                sessions_by_day=sessions_by_day, nudges=nudges, short_other=short_other,
                prompt_lens=prompt_lens, sessions=sessions, interrupts=interrupts,
                rejects=rejects, compactions=compactions, n_files=n_files)


def summarize(r):
    models = {}
    total_units = 0.0
    for m, t in r["tok"].items():
        units = (t["in"] * W_IN + t["cache_read"] * W_CACHE_READ
                 + t["cache_write"] * W_CACHE_WRITE + t["out"] * W_OUT)
        total_units += units
        reads = t["cache_read"]
        denom = reads + t["in"] + t["cache_write"]
        models[m] = {
            "messages": r["msgs_by_model"][m],
            "input": t["in"], "output": t["out"],
            "cache_read": reads, "cache_write": t["cache_write"],
            "cache_hit_rate": round(reads / denom, 4) if denom else None,
            "input_equiv_units": round(units),
        }
    for m in models.values():
        m["share_of_spend"] = round(m["input_equiv_units"] / total_units, 4) if total_units else None
    lens = sorted(r["prompt_lens"])
    ses = r["sessions"]
    n_sess = len(ses)
    total_prompts = len(lens)
    nudge_total = sum(r["nudges"].values())
    days = sorted(r["sessions_by_day"])
    return {
        "schema": 1,
        "range": {"first_day": days[0] if days else None, "last_day": days[-1] if days else None},
        "scale": {"transcript_files": r["n_files"], "main_sessions": n_sess,
                  "human_prompts": total_prompts, "compactions": r["compactions"],
                  "interrupts": r["interrupts"], "permission_rejects": r["rejects"]},
        "spend": {"weights": {"input": W_IN, "cache_read": W_CACHE_READ,
                              "cache_write": W_CACHE_WRITE, "output": W_OUT},
                  "total_input_equiv_units": round(total_units), "by_model": models},
        "nudges": {"total": nudge_total,
                   "share_of_prompts": round(nudge_total / total_prompts, 4) if total_prompts else None,
                   "by_phrase": dict(r["nudges"].most_common(15)),
                   "other_short_prompts": r["short_other"]},
        "prompts": {"median_len": lens[len(lens) // 2] if lens else None,
                    "p90_len": lens[int(len(lens) * .9)] if lens else None},
        "sessions": {"one_prompt_or_less": sum(1 for s in ses if s["prompts"] <= 1),
                     "over_500_assistant_msgs": sum(1 for s in ses if s["assistant_msgs"] > 500),
                     "ctx_peak_over_500k": sum(1 for s in ses if s["ctx_peak"] > 500_000),
                     "median_assistant_msgs": sorted(s["assistant_msgs"] for s in ses)[n_sess // 2] if n_sess else None},
        "tools": dict(r["tools"].most_common(30)),
        "tool_errors": dict(r["tool_err"].most_common(15)),
        "skills": dict(r["skills"].most_common(20)),
        "slash_commands": dict(r["slash"].most_common(20)),
        "agents": {"types": dict(r["agent_types"].most_common(15)),
                   "model_param": dict(r["agent_models"])},
        "mcp_servers": dict(r["mcp"].most_common(15)),
    }


def to_md(s):
    L = ["# Claude usage audit - aggregate summary", ""]
    L.append(f"Range: {s['range']['first_day']} .. {s['range']['last_day']}")
    sc = s["scale"]
    L.append(f"Sessions: {sc['main_sessions']} | human prompts: {sc['human_prompts']} | "
             f"compactions: {sc['compactions']} | interrupts: {sc['interrupts']} | rejects: {sc['permission_rejects']}")
    L += ["", "## Spend by model (input-equivalent units)"]
    for m, v in sorted(s["spend"]["by_model"].items(), key=lambda kv: -kv[1]["input_equiv_units"]):
        L.append(f"- {m}: {v['input_equiv_units']:,} ({(v['share_of_spend'] or 0):.1%}), "
                 f"cache hit {(v['cache_hit_rate'] or 0):.1%}, {v['messages']:,} msgs")
    n = s["nudges"]
    L += ["", "## Nudges (user restarting paused work)",
          f"- total {n['total']} = {(n['share_of_prompts'] or 0):.1%} of prompts"]
    L += [f"  - \"{k}\": {v}" for k, v in n["by_phrase"].items()]
    se = s["sessions"]
    L += ["", "## Session shape",
          f"- one-prompt-or-less sessions: {se['one_prompt_or_less']}",
          f"- sessions over 500 assistant messages: {se['over_500_assistant_msgs']}",
          f"- sessions with context peak over 500k: {se['ctx_peak_over_500k']}"]
    L += ["", "## Top tools"] + [f"- {k}: {v}" for k, v in list(s["tools"].items())[:12]]
    return "\n".join(L) + "\n"


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=os.path.expanduser("~/.claude/projects"))
    ap.add_argument("--out", default="./audit-out")
    ap.add_argument("--since", default=None, help="only count activity from this date (YYYY-MM-DD)")
    a = ap.parse_args()
    if not os.path.isdir(a.root):
        sys.exit(f"No transcripts found at {a.root}")
    s = summarize(scan(a.root, a.since))
    os.makedirs(a.out, exist_ok=True)
    with open(os.path.join(a.out, "summary.json"), "w", encoding="utf-8") as f:
        json.dump(s, f, indent=2)
    with open(os.path.join(a.out, "summary.md"), "w", encoding="utf-8") as f:
        f.write(to_md(s))
    print(f"Wrote {a.out}/summary.json and summary.md - review them, then send summary.json only.")


if __name__ == "__main__":
    main()