#!/usr/bin/env python3
"""Enumerate agent sessions in a window and dump their human turns.

Usage: scan-sessions.py [--days N] [--since YYYY-MM-DD] [--out DIR]
                       [--tools claude,codex,pi,opencode]

Covers four tools:

  claude    ~/.claude/projects/<slug>/<uuid>.jsonl
  codex     ~/.codex/sessions/**/rollout-*.jsonl, plus the thread index in
            ~/.codex/state_*.sqlite when the rollout files are gone
  pi        ~/.pi/agent/sessions/<slug>/<ts>_<uuid>.jsonl
  opencode  ~/.local/share/opencode/opencode-stable.db

Writes three files to --out (default: cwd):

  sessions.json   one record per top-level session, oldest first
  userturns.txt   every human turn, grouped by session, for reading
  models.md       spend and model/effort breakdown, ready to paste

A "top-level" session is one with no isSidechain marker and at least one real
human turn, which drops subagent transcripts. Swarm runs still show up as many
sessions sharing one cwd and timestamp — the report should collapse those.

Cost is reported by pi and opencode themselves. For Claude Code and Codex it is
estimated from token counts and scripts/pricing.json, and marked "estimated" —
a subscription seat does not bill this, so read it as the API-equivalent price
of the work, not as an invoice. A model missing from pricing.json produces no
cost at all and is listed under "unpriced" so the gap is visible.
"""
import argparse
import collections
import datetime as dt
import glob
import json
import os
import re
import sqlite3
import sys

SKIP_PREFIXES = (
    "<local-command", "<command-", "<task-notification", "<system-reminder",
    "Caveat:", "Base directory for this skill:", "Stop hook feedback:",
)

# Crude on purpose: a hit means the user pushed back or repeated themselves,
# which points at a session worth reading. It is not a quality score.
REDO = re.compile(
    r"\b(no,|nope|wrong|that'?s not|not what|still (broken|failing|wrong|there)|"
    r"again|revert|undo|as i said|i said|already (said|told)|stop |don'?t )",
    re.I,
)

PRICING = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pricing.json")

DATED = re.compile(r"-\d{8}$")


def norm_model(name):
    """Canonical model id, or None for a non-model.

    Claude Code stamps `<synthetic>` on messages it generated locally (API
    errors, interrupts) — those are not a model and must not appear in a spend
    table. Dated aliases like `claude-haiku-4-5-20251001` are the same model as
    the undated id the pricing table uses.
    """
    if not name or name.startswith("<"):
        return None
    return DATED.sub("", name)


def load_pricing(path=PRICING):
    try:
        with open(path) as fh:
            return json.load(fh)
    except OSError:
        print(f"no pricing table at {path} — costs will be blank",
              file=sys.stderr)
        return {"models": {}, "cache_write_multiplier": 1.25,
                "cache_read_multiplier": 0.1}


def blank():
    return {"input": 0, "output": 0, "cache_read": 0, "cache_write": 0,
            "reasoning": 0}


def add(dst, src):
    for k, v in src.items():
        dst[k] = dst.get(k, 0) + v


def estimate_cost(models, tokens, pricing):
    """USD for a session, split over the models it used.

    Token counts are per session, not per model, so a session that switched
    models mid-way is apportioned by assistant-message share. That is an
    approximation and only matters for mixed sessions, which are rare.
    """
    total_msgs = sum(models.values()) or 1
    cost, unpriced = 0.0, []
    cw = pricing.get("cache_write_multiplier", 1.25)
    cr = pricing.get("cache_read_multiplier", 0.1)
    for model, n in models.items():
        rate = pricing["models"].get(model)
        if not rate:
            unpriced.append(model)
            continue
        share = n / total_msgs
        cost += share * (
            tokens["input"] * rate["input"]
            + tokens["output"] * rate["output"]
            + tokens["cache_write"] * rate["input"] * cw
            + tokens["cache_read"] * rate["input"] * cr
        ) / 1e6
    return round(cost, 4), unpriced


def rec(tool, ts, cwd, file, mb=0.0, lines=0):
    return {"tool": tool, "ts": ts or "", "end": ts or "", "cwd": cwd or "",
            "mb": mb, "lines": lines, "models": {}, "efforts": {},
            "tokens": blank(), "cost_usd": None, "cost_source": None,
            "assistant_msgs": 0, "tool_errors": 0, "turns": [], "file": file}


def flatten(content):
    """Content list or string -> plain text of its text blocks."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(
            b.get("text", "") for b in content
            if isinstance(b, dict) and b.get("type") == "text"
        )
    return ""


def clean(text):
    text = " ".join((text or "").split())
    if not text or text.startswith(SKIP_PREFIXES):
        return None
    if "This session is being continued" in text[:60]:
        return None
    return text


# --- Claude Code -----------------------------------------------------------

def scan_claude(root, cutoff):
    for f in glob.glob(os.path.join(root, "*", "*.jsonl")):
        try:
            st = os.stat(f)
        except OSError:
            continue
        if st.st_mtime < cutoff:
            continue
        r = rec("claude", None, None, f, round(st.st_size / 1048576, 1))
        sidechain = False
        try:
            fh = open(f, errors="replace")
        except OSError:
            continue
        with fh:
            for line in fh:
                r["lines"] += 1
                try:
                    d = json.loads(line)
                except ValueError:
                    continue
                if d.get("isSidechain"):
                    sidechain = True
                    break
                ts = d.get("timestamp")
                if ts:
                    if not r["ts"]:
                        r["ts"] = ts
                    r["end"] = ts
                if not r["cwd"] and d.get("cwd"):
                    r["cwd"] = d["cwd"]
                m = d.get("message") or {}
                if d.get("type") == "assistant":
                    r["assistant_msgs"] += 1
                    model = norm_model(m.get("model"))
                    if model:
                        r["models"][model] = r["models"].get(model, 0) + 1
                    if d.get("effort"):
                        e = d["effort"]
                        r["efforts"][e] = r["efforts"].get(e, 0) + 1
                    u = m.get("usage") or {}
                    add(r["tokens"], {
                        "input": u.get("input_tokens", 0),
                        "output": u.get("output_tokens", 0),
                        "cache_read": u.get("cache_read_input_tokens", 0),
                        "cache_write": u.get("cache_creation_input_tokens", 0),
                    })
                elif d.get("type") == "user":
                    c = m.get("content")
                    if isinstance(c, list):
                        for b in c:
                            if isinstance(b, dict) and b.get("is_error"):
                                r["tool_errors"] += 1
                    t = clean(flatten(c))
                    if t:
                        r["turns"].append(t)
        if sidechain or not r["turns"]:
            continue
        if not r["cwd"]:
            r["cwd"] = os.path.basename(os.path.dirname(f))
        yield r


# --- pi --------------------------------------------------------------------

def scan_pi(root, cutoff):
    for f in glob.glob(os.path.join(root, "*", "*.jsonl")):
        try:
            st = os.stat(f)
        except OSError:
            continue
        if st.st_mtime < cutoff:
            continue
        r = rec("pi", None, None, f, round(st.st_size / 1048576, 1))
        cost, effort = 0.0, None
        try:
            fh = open(f, errors="replace")
        except OSError:
            continue
        with fh:
            for line in fh:
                r["lines"] += 1
                try:
                    d = json.loads(line)
                except ValueError:
                    continue
                ts = d.get("timestamp")
                if ts:
                    if not r["ts"]:
                        r["ts"] = ts
                    r["end"] = ts
                kind = d.get("type")
                if kind == "session" and d.get("cwd"):
                    r["cwd"] = r["cwd"] or d["cwd"]
                elif kind == "thinking_level_change":
                    effort = d.get("thinkingLevel")
                elif kind == "message":
                    m = d.get("message") or {}
                    role = m.get("role")
                    if role == "assistant":
                        r["assistant_msgs"] += 1
                        model = norm_model(m.get("model"))
                        if model:
                            r["models"][model] = r["models"].get(model, 0) + 1
                        if effort:
                            r["efforts"][effort] = r["efforts"].get(effort, 0) + 1
                        u = m.get("usage") or {}
                        add(r["tokens"], {
                            "input": u.get("input", 0),
                            "output": u.get("output", 0),
                            "cache_read": u.get("cacheRead", 0),
                            "cache_write": u.get("cacheWrite", 0),
                            "reasoning": u.get("reasoning", 0),
                        })
                        cost += ((u.get("cost") or {}).get("total") or 0)
                    elif role == "toolResult":
                        if m.get("isError"):
                            r["tool_errors"] += 1
                    elif role == "user":
                        t = clean(flatten(m.get("content")))
                        if t:
                            r["turns"].append(t)
        if not r["turns"]:
            continue
        r["cost_usd"] = round(cost, 4)
        r["cost_source"] = "reported"
        if not r["cwd"]:
            r["cwd"] = os.path.basename(os.path.dirname(f))
        yield r


# --- Codex -----------------------------------------------------------------

def scan_codex_rollouts(root, cutoff):
    """Rollout transcripts. Codex has changed this layout more than once, so
    every field here is read defensively and a miss costs a blank column, not
    a crash."""
    for f in glob.glob(os.path.join(root, "**", "*.jsonl"), recursive=True):
        try:
            st = os.stat(f)
        except OSError:
            continue
        if st.st_mtime < cutoff:
            continue
        r = rec("codex", None, None, f, round(st.st_size / 1048576, 1))
        try:
            fh = open(f, errors="replace")
        except OSError:
            continue
        with fh:
            for line in fh:
                r["lines"] += 1
                try:
                    d = json.loads(line)
                except ValueError:
                    continue
                ts = d.get("timestamp")
                if ts:
                    if not r["ts"]:
                        r["ts"] = ts
                    r["end"] = ts
                p = d.get("payload") if isinstance(d.get("payload"), dict) else d
                if p.get("cwd") and not r["cwd"]:
                    r["cwd"] = p["cwd"]
                model = norm_model(
                    p.get("model") or (p.get("turn_context") or {}).get("model"))
                if model:
                    r["models"][model] = r["models"].get(model, 0) + 1
                eff = (p.get("effort") or p.get("reasoning_effort")
                       or (p.get("turn_context") or {}).get("effort"))
                if eff:
                    r["efforts"][eff] = r["efforts"].get(eff, 0) + 1
                info = p.get("info") or {}
                usage = (info.get("last_token_usage") or info.get("total_token_usage")
                         or p.get("usage"))
                if isinstance(usage, dict):
                    add(r["tokens"], {
                        "input": usage.get("input_tokens", 0),
                        "output": usage.get("output_tokens", 0),
                        "cache_read": usage.get("cached_input_tokens", 0),
                        "reasoning": usage.get("reasoning_output_tokens", 0),
                    })
                if p.get("type") == "message" and p.get("role") == "user":
                    t = clean(flatten(p.get("content")))
                    if t:
                        r["turns"].append(t)
                elif p.get("role") == "assistant":
                    r["assistant_msgs"] += 1
        if not r["turns"]:
            continue
        yield r


def scan_codex_threads(home, cutoff, seen_paths):
    """Fallback index: threads Codex recorded in sqlite whose rollout file is
    gone or unparsed. Gives model, effort, tokens and the first user message,
    but no full turn list — enough to keep the session from vanishing from the
    week."""
    for db in sorted(glob.glob(os.path.join(home, "state_*.sqlite"))):
        try:
            con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
            rows = con.execute(
                "select rollout_path, created_at, updated_at, cwd, model, "
                "reasoning_effort, tokens_used, first_user_message, title "
                "from threads where updated_at >= ?", (int(cutoff),)
            ).fetchall()
            con.close()
        except sqlite3.Error as e:
            print(f"codex sqlite {db}: {e}", file=sys.stderr)
            continue
        for (path, created, updated, cwd, model, eff, tokens, first, title) in rows:
            if path and os.path.abspath(path) in seen_paths:
                continue
            r = rec("codex", dt.datetime.fromtimestamp(created).isoformat(),
                    cwd, path or db)
            r["end"] = dt.datetime.fromtimestamp(updated).isoformat()
            if model:
                r["models"][model] = 1
            if eff:
                r["efforts"][eff] = 1
            r["tokens"]["input"] = tokens or 0
            r["turns"] = [clean(first or title) or "(no user message recorded)"]
            r["partial"] = "sqlite index only — rollout transcript not read"
            yield r


# --- opencode --------------------------------------------------------------

def scan_opencode(db_path, cutoff):
    if not os.path.exists(db_path):
        return
    try:
        con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        con.row_factory = sqlite3.Row
        sessions = con.execute(
            "select * from session where time_updated >= ?",
            (int(cutoff * 1000),)
        ).fetchall()
    except sqlite3.Error as e:
        print(f"opencode sqlite: {e}", file=sys.stderr)
        return
    for s in sessions:
        r = rec("opencode",
                dt.datetime.fromtimestamp(s["time_created"] / 1000).isoformat(),
                s["directory"], f"{db_path}#{s['id']}")
        r["end"] = dt.datetime.fromtimestamp(s["time_updated"] / 1000).isoformat()
        r["title"] = s["title"]
        r["tokens"] = {
            "input": s["tokens_input"], "output": s["tokens_output"],
            "cache_read": s["tokens_cache_read"],
            "cache_write": s["tokens_cache_write"],
            "reasoning": s["tokens_reasoning"],
        }
        r["cost_usd"] = round(s["cost"] or 0, 4)
        r["cost_source"] = "reported"
        msgs = con.execute(
            "select id, data from message where session_id = ? "
            "order by time_created", (s["id"],)
        ).fetchall()
        user_ids = []
        for m in msgs:
            try:
                d = json.loads(m["data"])
            except ValueError:
                continue
            if d.get("role") == "assistant":
                r["assistant_msgs"] += 1
                model = norm_model(d.get("modelID"))
                if model:
                    r["models"][model] = r["models"].get(model, 0) + 1
                # opencode calls the effort level a model "variant".
                v = d.get("variant")
                if v:
                    r["efforts"][v] = r["efforts"].get(v, 0) + 1
            elif d.get("role") == "user":
                user_ids.append(m["id"])
        r["lines"] = len(msgs)
        for mid in user_ids:
            parts = con.execute(
                "select data from part where message_id = ? order by id", (mid,)
            ).fetchall()
            text = " ".join(
                json.loads(p["data"]).get("text", "")
                for p in parts
                if json.loads(p["data"]).get("type") == "text"
            )
            t = clean(text)
            # opencode asks the model to title the session through the same
            # message table; that prompt is not a human turn.
            if t and "Generate a concise 3 to 5 word title" not in t:
                r["turns"].append(t)
        errs = con.execute(
            "select count(*) from part where session_id = ? and "
            "json_extract(data,'$.state.status') = 'error'", (s["id"],)
        ).fetchone()
        r["tool_errors"] = errs[0] if errs else 0
        if r["turns"]:
            yield r
    con.close()


# --- report ----------------------------------------------------------------

def md_table(rows, right=()):
    w = [max(len(r[i]) for r in rows) for i in range(len(rows[0]))]
    def line(cells):
        return "| " + " | ".join(
            c.rjust(w[i]) if i in right else c.ljust(w[i])
            for i, c in enumerate(cells)) + " |"
    sep = "|" + "|".join(
        ("-" * (w[i] + 1) + ":") if i in right else ("-" * (w[i] + 2))
        for i in range(len(w))) + "|"
    return "\n".join([line(rows[0]), sep] + [line(r) for r in rows[1:]])


def top(counter, n=2):
    return ", ".join(f"{k}" for k, _ in
                     sorted(counter.items(), key=lambda x: -x[1])[:n]) or "—"


def breakdown(sessions):
    """Spend and friction per (tool, model, effort)."""
    agg = collections.defaultdict(lambda: {
        "sessions": 0, "turns": 0, "redo": 0, "errors": 0,
        "out": 0, "cost": 0.0, "priced": 0,
    })
    for s in sessions:
        model = top(s["models"], 1)
        effort = top(s["efforts"], 1)
        a = agg[(s["tool"], model, effort)]
        a["sessions"] += 1
        a["turns"] += len(s["turns"])
        a["redo"] += sum(1 for t in s["turns"] if REDO.search(t))
        a["errors"] += s["tool_errors"]
        a["out"] += s["tokens"]["output"]
        if s["cost_usd"] is not None:
            a["cost"] += s["cost_usd"]
            a["priced"] += 1
    rows = [["tool", "model", "effort", "sess", "turns", "push-back",
             "tool err", "out tok", "$", "$/turn"]]
    for (tool, model, effort), a in sorted(
            agg.items(), key=lambda x: -x[1]["cost"]):
        per = a["cost"] / a["turns"] if a["turns"] and a["cost"] else 0
        rows.append([
            tool, model, effort, str(a["sessions"]), str(a["turns"]),
            str(a["redo"]), str(a["errors"]), f"{a['out']:,}",
            f"{a['cost']:.2f}" if a["cost"] else "—",
            f"{per:.3f}" if per else "—",
        ])
    return md_table(rows, right=set(range(3, 10)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--days", type=int, default=7)
    ap.add_argument("--since")
    ap.add_argument("--out", default=".")
    ap.add_argument("--tools", default="claude,codex,pi,opencode")
    ap.add_argument("--claude-root",
                    default=os.path.expanduser("~/.claude/projects"))
    ap.add_argument("--codex-home", default=os.path.expanduser("~/.codex"))
    ap.add_argument("--pi-root",
                    default=os.path.expanduser("~/.pi/agent/sessions"))
    ap.add_argument("--opencode-db", default=os.path.expanduser(
        "~/.local/share/opencode/opencode-stable.db"))
    a = ap.parse_args()

    if a.since:
        cutoff = dt.datetime.fromisoformat(a.since).timestamp()
    else:
        cutoff = (dt.datetime.now() - dt.timedelta(days=a.days)).timestamp()
    want = {t.strip() for t in a.tools.split(",") if t.strip()}
    pricing = load_pricing()

    sessions = []
    if "claude" in want and os.path.isdir(a.claude_root):
        sessions += list(scan_claude(a.claude_root, cutoff))
    if "pi" in want and os.path.isdir(a.pi_root):
        sessions += list(scan_pi(a.pi_root, cutoff))
    if "codex" in want:
        rollouts = list(scan_codex_rollouts(
            os.path.join(a.codex_home, "sessions"), cutoff))
        sessions += rollouts
        sessions += list(scan_codex_threads(
            a.codex_home, cutoff,
            {os.path.abspath(r["file"]) for r in rollouts}))
    if "opencode" in want:
        sessions += list(scan_opencode(a.opencode_db, cutoff))

    unpriced = set()
    for s in sessions:
        if s["cost_usd"] is None:
            s["cost_usd"], miss = estimate_cost(s["models"], s["tokens"], pricing)
            s["cost_source"] = "estimated"
            unpriced.update(miss)
            if not s["cost_usd"]:
                s["cost_usd"] = None
                s["cost_source"] = None
    sessions.sort(key=lambda x: x["ts"])

    os.makedirs(a.out, exist_ok=True)
    with open(os.path.join(a.out, "sessions.json"), "w") as fh:
        json.dump([{k: v for k, v in s.items() if k != "turns"}
                   for s in sessions], fh, indent=1)

    with open(os.path.join(a.out, "userturns.txt"), "w") as fh:
        for s in sessions:
            cost = (f"${s['cost_usd']:.2f}"
                    f"{'~' if s['cost_source'] == 'estimated' else ''}"
                    if s["cost_usd"] else "$?")
            fh.write(f"\n===== {s['ts'][:16]} [{s['tool']}] {s['cwd']} "
                     f"({s['mb']}MB, {top(s['models'])}, "
                     f"effort {top(s['efforts'])}, {cost}) =====\n")
            if s.get("partial"):
                fh.write(f"  ({s['partial']})\n")
            for t in s["turns"]:
                fh.write("- " + t[:500] + "\n")

    report = breakdown(sessions)
    with open(os.path.join(a.out, "models.md"), "w") as fh:
        fh.write("# Model, effort and spend, past window\n\n")
        fh.write(report + "\n\n")
        fh.write("`push-back` counts human turns matching a crude "
                 "correction regex — a pointer to sessions worth reading, "
                 "not a quality score.\n")
        fh.write("Cost is reported by pi and opencode, estimated from "
                 "pricing.json for Claude Code and Codex.\n")
        if unpriced:
            fh.write("\nUnpriced models (no cost counted): "
                     + ", ".join(sorted(unpriced)) + "\n")

    by_tool = collections.Counter(s["tool"] for s in sessions)
    print(f"{len(sessions)} top-level sessions "
          f"({', '.join(f'{v} {k}' for k, v in by_tool.most_common())}), "
          f"{sum(len(s['turns']) for s in sessions)} human turns, "
          f"{sum(s['mb'] for s in sessions):.0f}MB")
    print()
    print(report)
    if unpriced:
        print("\nunpriced models (add them to scripts/pricing.json): "
              + ", ".join(sorted(unpriced)))

    by_cwd = collections.Counter(s["cwd"] for s in sessions)
    swarms = {k: v for k, v in by_cwd.items() if v > 20}
    if swarms:
        print("\nlikely swarm runs (collapse these to one line in the report):")
        for k, v in sorted(swarms.items(), key=lambda x: -x[1]):
            print(f"  {v:4} sessions  {k}")
    print(f"\nwrote {a.out}/sessions.json, {a.out}/userturns.txt, "
          f"{a.out}/models.md")


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