#!/usr/bin/env python3
"""Deterministic aggregator for a single Claude Code session.

Reads the main transcript JSONL for one session plus its subagent
sidechain logs and its Workflow-substrate run logs (workflows/wf_*.json
+ subagents/workflows/<runId>/agent-*.jsonl), and emits a compact JSON
summary on stdout that the post-mortem SKILL interprets into a written
report.

Why a script and not skill prose: transcripts run to >1 MB and a
session can spawn dozens of subagent logs. Token accounting has two
traps that only deterministic code gets right every time:

  1. Streamed assistant lines repeat the SAME requestId with the SAME
     usage object. Counting per-line double-counts tokens. We dedup
     per requestId (max per field within a request, then sum across
     requests).
  2. The transcript carries NO durationMs (always null), so wall-clock
     is derived from timestamp deltas. The numbers reported here are
     raw token counts, not a dollar estimate — different token classes
     (input / output / cache) are reported separately rather than
     collapsed into a single priced figure.

Stdlib only. No network. Read-only against ~/.claude.
"""

import argparse
import glob
import json
import os
import re
import sys
from datetime import datetime, timezone

_TOK_ALT = "DONE_WITH_CONCERNS|NEEDS_CONTEXT|BLOCKED|PARTIAL|DONE"  # longest-first
# Authoritative signal: the implement end-report leads with a structured
# "Status:  DONE" line (optionally wrapped in markdown bold or a list
# marker). Anchoring here is what stops the scan from flipping a shipped
# DONE run to BLOCKED on the strength of the report's own template lines
# (`BLOCKED file:  BLOCKED.md`, `Blocked detail:`) or prose narrating a
# surmounted blocker. See issue #3.
STATUS_LINE_RE = re.compile(
    r"^[\s>*#\-]*\**\s*status\s*\**\s*[:\-]\s*\**\s*(" + _TOK_ALT + r")\b",
    re.IGNORECASE | re.MULTILINE)
# Fallback for agents that emit a bare status token as their final line
# (no "Status:" prefix). Requires the token to BE the line — excludes
# substrings like "BLOCKED.md" or "BLOCKED file:" buried in prose.
STATUS_STANDALONE_RE = re.compile(
    r"^[\s>*#\-]*\**\s*(" + _TOK_ALT + r")\**\s*$",
    re.IGNORECASE | re.MULTILINE)


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


def empty_tokens():
    return {"input": 0, "output": 0, "cache_creation": 0, "cache_read": 0,
            "cache_creation_5m": 0, "cache_creation_1h": 0}


def token_total(tokens):
    """Raw count of every token the session moved. cache_creation_5m/1h are
    a sub-split of cache_creation, so they are excluded to avoid double
    counting."""
    return (tokens["input"] + tokens["output"]
            + tokens["cache_creation"] + tokens["cache_read"])


def add_tokens(acc, usage):
    """Add one request's deduped usage into an accumulator."""
    acc["input"] += usage.get("input_tokens", 0) or 0
    acc["output"] += usage.get("output_tokens", 0) or 0
    cc = usage.get("cache_creation_input_tokens", 0) or 0
    acc["cache_creation"] += cc
    acc["cache_read"] += usage.get("cache_read_input_tokens", 0) or 0
    detail = usage.get("cache_creation") or {}
    f5 = detail.get("ephemeral_5m_input_tokens")
    f1 = detail.get("ephemeral_1h_input_tokens")
    if f5 is None and f1 is None:
        acc["cache_creation_5m"] += cc  # no split available -> assume 5m
    else:
        acc["cache_creation_5m"] += f5 or 0
        acc["cache_creation_1h"] += f1 or 0


def iter_lines(path):
    with open(path, "r", encoding="utf-8", errors="replace") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                continue


def aggregate_transcript(path):
    """Return token totals, model mix, and per-request count for one JSONL
    transcript (main or subagent). Dedups usage by requestId."""
    # requestId -> {field: max seen}, plus model per request.
    req_usage = {}
    req_model = {}
    for rec in iter_lines(path):
        if rec.get("type") != "assistant":
            continue
        msg = rec.get("message") or {}
        usage = msg.get("usage")
        if not usage:
            continue
        key = rec.get("requestId") or rec.get("uuid")
        model = msg.get("model")
        if model:
            req_model[key] = model
        cur = req_usage.setdefault(key, {})
        # Max-per-field guards against both identical-repeat and
        # cumulative-within-request streaming shapes.
        for f in ("input_tokens", "output_tokens",
                  "cache_creation_input_tokens", "cache_read_input_tokens"):
            v = usage.get(f, 0) or 0
            if v > cur.get(f, 0):
                cur[f] = v
        det = usage.get("cache_creation") or {}
        cd = cur.setdefault("cache_creation", {})
        for f in ("ephemeral_5m_input_tokens", "ephemeral_1h_input_tokens"):
            v = det.get(f, 0) or 0
            if v > cd.get(f, 0):
                cd[f] = v
        st = usage.get("server_tool_use") or {}
        sc = cur.setdefault("_server", {})
        for f in ("web_search_requests", "web_fetch_requests"):
            v = st.get(f, 0) or 0
            if v > sc.get(f, 0):
                sc[f] = v

    tokens = empty_tokens()
    server = {"web_search": 0, "web_fetch": 0}
    models = {}
    for key, usage in req_usage.items():
        add_tokens(tokens, usage)
        model = req_model.get(key)
        models[model] = models.get(model, 0) + 1
        sv = usage.get("_server", {})
        server["web_search"] += sv.get("web_search_requests", 0)
        server["web_fetch"] += sv.get("web_fetch_requests", 0)

    return {
        "requests": len(req_usage),
        "tokens": tokens,
        "total_tokens": token_total(tokens),
        "models": models,
        "server_tools": server,
    }


def cache_hit_ratio(tokens):
    denom = tokens["input"] + tokens["cache_read"] + tokens["cache_creation"]
    return round(tokens["cache_read"] / denom, 4) if denom else None


def analyze_main(path):
    """Tool mix, error rate, skills, slash commands, prompts, wall-clock."""
    tool_calls = {}
    error_total = 0
    interrupted = 0
    skills = []
    slash = []
    tasks = []
    user_prompts = 0
    first_ts = last_ts = None
    cwd = version = None

    for rec in iter_lines(path):
        ts = parse_ts(rec.get("timestamp"))
        if ts:
            if first_ts is None or ts < first_ts:
                first_ts = ts
            if last_ts is None or ts > last_ts:
                last_ts = ts
        cwd = rec.get("cwd") or cwd
        version = rec.get("version") or version

        typ = rec.get("type")
        msg = rec.get("message") or {}
        content = msg.get("content")

        if typ == "assistant" and isinstance(content, list):
            for block in content:
                if block.get("type") != "tool_use":
                    continue
                name = block.get("name", "?")
                tool_calls[name] = tool_calls.get(name, 0) + 1
                if name == "Skill":
                    skills.append({"skill": (block.get("input") or {}).get("skill"),
                                   "ts": rec.get("timestamp")})
                elif name in ("Task", "Agent"):
                    inp = block.get("input") or {}
                    tasks.append({"subagent_type": inp.get("subagent_type"),
                                  "description": inp.get("description"),
                                  "ts": rec.get("timestamp")})

        if typ == "user":
            # tool_result blocks carry the authoritative is_error flag
            if isinstance(content, list):
                for block in content:
                    if block.get("type") == "tool_result":
                        if block.get("is_error"):
                            error_total += 1
                    elif block.get("type") == "text":
                        for m in re.findall(r"<command-name>(.*?)</command-name>",
                                            block.get("text", "")):
                            slash.append({"cmd": m, "ts": rec.get("timestamp")})
            elif isinstance(content, str):
                for m in re.findall(r"<command-name>(.*?)</command-name>", content):
                    slash.append({"cmd": m, "ts": rec.get("timestamp")})
            # genuine user turns (not tool results, not slash plumbing)
            is_tool_result = (isinstance(content, list)
                              and any(b.get("type") == "tool_result" for b in content))
            if not is_tool_result and rec.get("toolUseResult") is None:
                user_prompts += 1
            tur = rec.get("toolUseResult")
            if isinstance(tur, dict) and tur.get("interrupted"):
                interrupted += 1

    return {
        "tool_calls": tool_calls,
        "total_tool_calls": sum(tool_calls.values()),
        "tool_error_total": error_total,
        "interrupted": interrupted,
        "skills_invoked": skills,
        "slash_commands": slash,
        "tasks_dispatched": tasks,
        "user_prompts": user_prompts,
        "wall_clock_sec": round((last_ts - first_ts).total_seconds(), 1)
        if first_ts and last_ts else None,
        "started": first_ts.isoformat() if first_ts else None,
        "ended": last_ts.isoformat() if last_ts else None,
        "cwd": cwd,
        "version": version,
    }


def subagent_terminal_status(path):
    last_text = ""
    for rec in iter_lines(path):
        if rec.get("type") != "assistant":
            continue
        content = (rec.get("message") or {}).get("content")
        if isinstance(content, list):
            for block in content:
                if block.get("type") == "text" and block.get("text", "").strip():
                    last_text = block["text"]
    # Primary: the structured "Status:" line. Take the LAST one so a
    # restated final status wins over any earlier mention.
    structured = STATUS_LINE_RE.findall(last_text)
    if structured:
        return structured[-1].upper()
    # Fallback: a status token standing alone on its own line.
    standalone = STATUS_STANDALONE_RE.findall(last_text)
    if standalone:
        return standalone[-1].upper()
    return "unknown"


def analyze_workflows(session_dir):
    """Aggregate Workflow-substrate runs: workflows/wf_*.json (per-agent
    summary: stage label, model, live token counter) joined with the raw
    per-agent transcripts under subagents/workflows/<runId>/agent-*.jsonl.

    The raw transcripts are authoritative for token counts (deduped by
    requestId, same as every other transcript here); the wf_*.json
    'tokens' counters under-track vs. the deduped raw usage (~11% in the
    206-run corpus that motivated this path, issue #28) and are reported
    only as a cross-check. Stage = the label up to the first ':'
    ('impl:3' -> 'impl'), so loop stages aggregate across tasks/rounds."""
    runs = []
    for wf_path in sorted(glob.glob(os.path.join(session_dir, "workflows",
                                                 "wf_*.json"))):
        try:
            with open(wf_path) as fh:
                data = json.load(fh)
        except (json.JSONDecodeError, OSError):
            continue
        run_id = data.get("runId") or os.path.basename(wf_path)[:-len(".json")]
        raw_dir = os.path.join(session_dir, "subagents", "workflows", run_id)
        meta_by_agent = {}
        for entry in data.get("workflowProgress") or []:
            if entry.get("type") != "workflow_agent" or not entry.get("agentId"):
                continue
            meta_by_agent[entry["agentId"]] = entry

        run_tokens = empty_tokens()
        stages = {}
        agent_count = 0
        untracked = 0
        for jl in sorted(glob.glob(os.path.join(raw_dir, "agent-*.jsonl"))):
            aid = os.path.basename(jl)[len("agent-"):-len(".jsonl")]
            agg = aggregate_transcript(jl)
            entry = meta_by_agent.get(aid)
            if entry is None:
                # Transcript with no progress entry (retried/dropped
                # attempt): keep it counted, grouped under '?', never
                # silently folded into a named stage.
                untracked += 1
            label = (entry or {}).get("label") or "?"
            stage = label.split(":", 1)[0]
            st = stages.setdefault(stage, {"count": 0, "tokens": empty_tokens(),
                                           "models": {}})
            st["count"] += 1
            agent_count += 1
            for k in run_tokens:
                run_tokens[k] += agg["tokens"][k]
                st["tokens"][k] += agg["tokens"][k]
            fallback_model = (entry or {}).get("model")
            for m, n in agg["models"].items():
                key = m or fallback_model
                st["models"][key] = st["models"].get(key, 0) + n

        for st in stages.values():
            st["total_tokens"] = token_total(st["tokens"])

        runs.append({
            "run_id": run_id,
            "name": data.get("workflowName"),
            "status": data.get("status"),
            "agents": agent_count,
            "agents_untracked": untracked,
            "reported_total_tokens": data.get("totalTokens"),
            "tokens": run_tokens,
            "total_tokens": token_total(run_tokens),
            "stages": stages,
        })
    return runs


def analyze_subagents(session_dir):
    out = []
    sub_dir = os.path.join(session_dir, "subagents")
    for jl in sorted(glob.glob(os.path.join(sub_dir, "agent-*.jsonl"))):
        meta = {}
        mp = jl.replace(".jsonl", ".meta.json")
        if os.path.exists(mp):
            try:
                with open(mp) as fh:
                    meta = json.load(fh)
            except (json.JSONDecodeError, OSError):
                pass
        agg = aggregate_transcript(jl)
        out.append({
            "agent_id": os.path.basename(jl)[len("agent-"):-len(".jsonl")],
            "agent_type": meta.get("agentType"),
            "description": meta.get("description"),
            "dispatched_by_tool": meta.get("toolUseId"),
            "models": agg["models"],
            "requests": agg["requests"],
            "tokens": agg["tokens"],
            "total_tokens": agg["total_tokens"],
            "terminal_status": subagent_terminal_status(jl),
        })
    return out


def resolve_session(args):
    """Return (main_jsonl_path, session_dir_or_None)."""
    if args.file:
        path = os.path.abspath(args.file)
        return path, path[:-len(".jsonl")] if path.endswith(".jsonl") else None

    cwd = args.cwd or os.getcwd()
    proj_dir = args.project_dir or os.path.join(
        os.path.expanduser("~/.claude/projects"), cwd.replace("/", "-"))
    if not os.path.isdir(proj_dir):
        raise SystemExit(f"no project log dir for cwd {cwd}: {proj_dir} not found")

    if args.session:
        path = os.path.join(proj_dir, args.session + ".jsonl")
        if not os.path.exists(path):
            raise SystemExit(f"session {args.session} not found in {proj_dir}")
        return path, os.path.join(proj_dir, args.session)

    # default: newest .jsonl in the project dir = current/last session
    candidates = sorted(glob.glob(os.path.join(proj_dir, "*.jsonl")),
                        key=os.path.getmtime, reverse=True)
    if not candidates:
        raise SystemExit(f"no session transcripts in {proj_dir}")
    path = candidates[0]
    return path, path[:-len(".jsonl")]


def main():
    ap = argparse.ArgumentParser(description="Single-session post-mortem aggregator.")
    ap.add_argument("--session", help="explicit session id (UUID)")
    ap.add_argument("--file", help="explicit path to a main transcript .jsonl")
    ap.add_argument("--cwd", help="project cwd to resolve logs for (default: $PWD)")
    ap.add_argument("--project-dir", help="explicit ~/.claude/projects/<slug> dir")
    args = ap.parse_args()

    main_path, session_dir = resolve_session(args)
    session_id = os.path.basename(main_path)[:-len(".jsonl")]

    main_agg = aggregate_transcript(main_path)
    main_meta = analyze_main(main_path)

    subagents = []
    workflows = []
    if session_dir and os.path.isdir(session_dir):
        subagents = analyze_subagents(session_dir)
        workflows = analyze_workflows(session_dir)

    sub_tokens = empty_tokens()
    for s in subagents:
        for k in sub_tokens:
            sub_tokens[k] += s["tokens"][k]

    wf_tokens = empty_tokens()
    wf_agents = 0
    for w in workflows:
        wf_agents += w["agents"]
        for k in wf_tokens:
            wf_tokens[k] += w["tokens"][k]

    total_tokens = empty_tokens()
    for k in total_tokens:
        total_tokens[k] = (main_agg["tokens"][k] + sub_tokens[k]
                           + wf_tokens[k])

    warnings = []
    is_active = False
    if main_meta["ended"]:
        ended = parse_ts(main_meta["ended"])
        if ended and (datetime.now(timezone.utc) - ended).total_seconds() < 120:
            is_active = True
            warnings.append("Session appears ACTIVE (last event <2 min ago); "
                            "totals are partial and will grow.")
    warnings.append("Token counts are raw usage; wall-clock is derived from "
                    "timestamp deltas (the transcript carries no durationMs).")

    report = {
        "session": {
            "id": session_id,
            "transcript": main_path,
            "cwd": main_meta["cwd"],
            "version": main_meta["version"],
            "started": main_meta["started"],
            "ended": main_meta["ended"],
            "wall_clock_sec": main_meta["wall_clock_sec"],
            "active": is_active,
        },
        "main": {
            "requests": main_agg["requests"],
            "models": main_agg["models"],
            "tokens": main_agg["tokens"],
            "total_tokens": main_agg["total_tokens"],
            "cache_hit_ratio": cache_hit_ratio(main_agg["tokens"]),
            "server_tools": main_agg["server_tools"],
            "user_prompts": main_meta["user_prompts"],
        },
        "tools": {
            "calls": main_meta["tool_calls"],
            "total": main_meta["total_tool_calls"],
            "error_total": main_meta["tool_error_total"],
            "error_ratio": round(main_meta["tool_error_total"]
                                 / main_meta["total_tool_calls"], 4)
            if main_meta["total_tool_calls"] else None,
            "interrupted": main_meta["interrupted"],
        },
        "skills_invoked": main_meta["skills_invoked"],
        "slash_commands": main_meta["slash_commands"],
        "tasks_dispatched": main_meta["tasks_dispatched"],
        "subagents": subagents,
        "subagent_totals": {
            "count": len(subagents),
            "tokens": sub_tokens,
            "total_tokens": token_total(sub_tokens),
        },
        "workflows": workflows,
        "workflow_totals": {
            "runs": len(workflows),
            "agents": wf_agents,
            "tokens": wf_tokens,
            "total_tokens": token_total(wf_tokens),
        },
        "totals": {
            "tokens": total_tokens,
            "total_tokens": token_total(total_tokens),
            "cache_hit_ratio": cache_hit_ratio(total_tokens),
        },
        "warnings": warnings,
    }
    json.dump(report, sys.stdout, indent=2)
    sys.stdout.write("\n")


if __name__ == "__main__":
    main()