#!/usr/bin/env python3
"""phase-timings.py prints how long each stage of a Claude Code workflow run took.

    python3 phase-timings.py RUN_DIR [RUN_DIR ...] [--gap 20] [--by phase|label] [--csv]

RUN_DIR is one run's transcript folder. Claude Code writes it at
    ~/.claude/projects/<project>/<session>/subagents/workflows/<runId>/
Give several folders to count one piece of work that took several runs.

It reads these files, all written by Claude Code:
    journal.jsonl          one "started" line and one "result" line per agent
    agent-<id>.jsonl       the agent's transcript, with a timestamp on every line
    agent-<id>.meta.json   the agent's label and phase, when the journal lacks them
    <session>/workflows/<runId>.json
                           the run record, for older runs whose journal has no
                           labels, and for the harness's own total of the run

How it counts:
    An agent's start and end are the first and last timestamps in its
    transcript, because the journal has no timestamps.
    A stage's wall time is the union of its agents' spans, so agents that ran
    side by side count once.
    Working time counts only the minutes when at least one agent was at work.
    It leaves out idle time between agents, and any silence longer than --gap
    minutes inside one agent, e.g. a stall on one tool call.
    Agent time adds up the working time of every agent.
    A gap is a stretch longer than --gap minutes when no agent was at work. It
    is called a usage limit when an agent logged a rate_limit error in the 10
    minutes before it began.

It prints only times, labels and counts. It never prints what an agent said,
so its output is safe to paste into a chat or a document.
"""
import argparse
import json
import os
import sys
from datetime import datetime, timedelta


def parse_time(text):
    return datetime.fromisoformat(text.replace("Z", "+00:00"))


def run_record(run_dir):
    """The harness's record of the run, kept beside the session's scripts."""
    run_dir = os.path.normpath(run_dir)
    session = os.path.dirname(os.path.dirname(os.path.dirname(run_dir)))
    path = os.path.join(session, "workflows", os.path.basename(run_dir) + ".json")
    if not os.path.exists(path):
        return None
    with open(path) as record:
        return json.load(record)


def read_agents(run_dir):
    """Return {agent_id: {label, phase, returned}} for every agent the journal started."""
    agents = {}
    with open(os.path.join(run_dir, "journal.jsonl")) as journal:
        for line in journal:
            if not line.strip():
                continue
            event = json.loads(line)
            if event.get("type") == "started":
                agents[event["agentId"]] = {"label": event.get("label"), "phase": event.get("phase"), "returned": False}
            elif event.get("type") == "result" and event.get("agentId") in agents:
                agents[event["agentId"]]["returned"] = True

    # Older journals carry only the agent id. Fill the label and the phase from
    # the agent's meta file, then from the run record.
    record = run_record(run_dir) or {}
    from_record = {entry.get("agentId"): entry for entry in record.get("workflowProgress") or [] if entry.get("type") == "workflow_agent"}
    for agent_id, info in agents.items():
        meta_path = os.path.join(run_dir, "agent-%s.meta.json" % agent_id)
        meta = {}
        if os.path.exists(meta_path):
            with open(meta_path) as meta_file:
                meta = json.load(meta_file)
        entry = from_record.get(agent_id, {})
        info["label"] = info["label"] or meta.get("description") or entry.get("label") or "?"
        info["phase"] = info["phase"] or meta.get("workflowPhase") or entry.get("phaseTitle") or "(no phase)"
    return agents


def read_transcript(run_dir, agent_id, gap):
    """Return (segments, limit_times) for one agent.

    A segment is a [start, end] stretch of steady work. A new segment begins
    when two lines are more than `gap` apart. limit_times are the times the
    agent logged a rate_limit error.
    """
    path = os.path.join(run_dir, "agent-%s.jsonl" % agent_id)
    if not os.path.exists(path):
        return [], []
    times, limits = [], []
    with open(path) as transcript:
        for line in transcript:
            try:
                entry = json.loads(line)
            except ValueError:
                continue
            if not entry.get("timestamp"):
                continue
            when = parse_time(entry["timestamp"])
            times.append(when)
            if entry.get("isApiErrorMessage") and entry.get("error") == "rate_limit":
                limits.append(when)
    times.sort()
    segments = []
    for when in times:
        if segments and when - segments[-1][1] <= gap:
            segments[-1][1] = when
        else:
            segments.append([when, when])
    return segments, limits


def union(intervals):
    merged = []
    for start, end in sorted(intervals):
        if merged and start <= merged[-1][1]:
            merged[-1][1] = max(merged[-1][1], end)
        else:
            merged.append([start, end])
    return merged


def total(intervals):
    return sum((end - start for start, end in intervals), timedelta())


def hm(delta):
    minutes = int(round(delta.total_seconds() / 60))
    return "%d:%02d" % (minutes // 60, minutes % 60)


def minutes(delta):
    return delta.total_seconds() / 60


def stage_of(info, by):
    if by == "label":
        # "build1:B3" gives "build1", and "reverify2:V1-laptop" gives "reverify2".
        return info["label"].split(":")[0]
    return info["phase"]


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("run_dirs", nargs="+", help="one or more run transcript folders")
    parser.add_argument("--gap", type=float, default=20, help="minutes of silence that count as a gap (default 20)")
    parser.add_argument("--by", choices=["phase", "label"], default="phase", help="group by the phase() title, or by the first part of the label")
    parser.add_argument("--csv", action="store_true", help="print CSV rows instead of a table")
    args = parser.parse_args()
    gap = timedelta(minutes=args.gap)

    for run_dir in args.run_dirs:
        if not os.path.exists(os.path.join(run_dir, "journal.jsonl")):
            sys.exit("%s has no journal.jsonl. Give a folder under <session>/subagents/workflows/." % run_dir)

    stages = {}
    all_segments, all_limits = [], []
    for run_dir in args.run_dirs:
        for agent_id, info in read_agents(run_dir).items():
            segments, limits = read_transcript(run_dir, agent_id, gap)
            all_limits += limits
            stage = stages.setdefault(stage_of(info, args.by), {"spans": [], "segments": [], "agents": 0, "lost": 0, "agent_time": timedelta()})
            stage["agents"] += 1
            stage["lost"] += 0 if info["returned"] else 1
            if segments:
                stage["spans"].append([segments[0][0], segments[-1][1]])
                stage["segments"] += segments
                stage["agent_time"] += total(segments)
                all_segments += segments

    if not all_segments:
        sys.exit("no agent transcripts with timestamps were found")
    rows = sorted(((name, s) for name, s in stages.items() if s["spans"]), key=lambda row: min(span[0] for span in row[1]["spans"]))
    names = [os.path.basename(os.path.normpath(d)) for d in args.run_dirs]

    if args.csv:
        print("runs,stage,agents,lost,start_utc,end_utc,wall_min,working_min,agent_min")
        for name, s in rows:
            spans = union(s["spans"])
            print("%s,%s,%d,%d,%s,%s,%.1f,%.1f,%.1f" % (
                " ".join(names), name, s["agents"], s["lost"],
                spans[0][0].strftime("%Y-%m-%dT%H:%M:%SZ"), spans[-1][1].strftime("%Y-%m-%dT%H:%M:%SZ"),
                minutes(total(spans)), minutes(total(union(s["segments"]))), minutes(s["agent_time"])))
        return

    print("runs: %s" % ", ".join(names))
    print("%-22s %6s %4s  %-11s  %-11s  %6s %7s %10s" % ("stage", "agents", "lost", "start UTC", "end UTC", "wall", "working", "agent time"))
    for name, s in rows:
        spans = union(s["spans"])
        print("%-22s %6d %4d  %-11s  %-11s  %6s %7s %10s" % (
            name[:22], s["agents"], s["lost"],
            spans[0][0].strftime("%m-%d %H:%M"), spans[-1][1].strftime("%m-%d %H:%M"),
            hm(total(spans)), hm(total(union(s["segments"]))), hm(s["agent_time"])))

    working = union(all_segments)
    print()
    print("elapsed, first start to last end   %s" % hm(working[-1][1] - working[0][0]))
    print("working, some agent at work        %s" % hm(total(working)))
    print("agent time, all agents added up    %s" % hm(sum((s["agent_time"] for s in stages.values()), timedelta())))
    gaps = [(a[1], b[0]) for a, b in zip(working, working[1:]) if b[0] - a[1] > gap]
    print("gaps over %g minutes: %d" % (args.gap, len(gaps)))
    for start, end in gaps:
        limit = any(start - timedelta(minutes=10) <= t <= end for t in all_limits)
        print("  %s to %s UTC  %6s  %s" % (start.strftime("%m-%d %H:%M"), end.strftime("%m-%d %H:%M"), hm(end - start),
                                        "usage limit" if limit else "no agent at work: a stop, a stall or a wait"))
    for run_dir, name in zip(args.run_dirs, names):
        record = run_record(run_dir)
        if record and record.get("durationMs") is not None:
            print("Claude Code's record of %s, last launch only: %s, %s, %s agents, %s tokens" % (
                name, record.get("workflowName"), hm(timedelta(milliseconds=record["durationMs"])),
                record.get("agentCount"), "{:,}".format(record.get("totalTokens") or 0)))


if __name__ == "__main__":
    main()