#!/usr/bin/env python3
"""Stop hook: post-run usage + warning report for completed workflow runs.

Fires every turn-end. Scans this session's workflows/ journals for runs that
have reached a terminal status since last seen (tracked in a per-session state
file) and surfaces, via systemMessage, for each newly-finalized run:
  - run line: runId, status glyph, wall-clock (journal durationMs), agent count
  - per-model token/cache breakdown (subagent transcripts, windowed to the run)
  - journal-authoritative total tokens
  - the run's accumulated/session .warnignore suppression report (result.warnSummary)

Numbers are honest about their seams: the total is the journal's; the per-model
split is transcript-derived and time-windowed, so the parts need not sum to the
total, and the window is exact only when one run owned its time slice.

Fail-open: any error exits 0 with no output. Never blocks the stop, never sets
decision/continue. Dedup via the state file makes the once-per-turn firing emit
each run exactly once.
"""
import sys, os, json, glob, tempfile, time

sys.dont_write_bytecode = True
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lib"))
try:
    from usage_agg import parse_files, render_parts, fmt_duration
except Exception:
    sys.exit(0)

GLYPH = {"completed": "✓", "halted": "⚠", "failed": "✗", "error": "✗", "killed": "✗"}


def load_json(path):
    try:
        with open(path) as f:
            return json.load(f)
    except Exception:
        return None


def is_terminal(j):
    return (
        isinstance(j, dict)
        and isinstance(j.get("status"), str)
        and j.get("status") != "running"
        and isinstance(j.get("durationMs"), (int, float))
    )


def render_run(j, subagent_files):
    run_id = j.get("runId", "?")
    status = j.get("status", "?")
    glyph = GLYPH.get(status, status)
    dur = fmt_duration(j.get("durationMs", 0) / 1000.0)
    agents = j.get("agentCount", 0)
    total = j.get("totalTokens", 0)
    start = j.get("startTime")
    window = None
    if isinstance(start, (int, float)):
        window = (start, start + j.get("durationMs", 0))
    per_model, _, _, _ = parse_files(subagent_files, window=window)
    lines = [f"run {run_id} {glyph} {dur}  {agents} agents"]
    lines += render_parts(per_model, color=False)
    try:
        lines.append(f"total {int(total):,} tok")
    except Exception:
        lines.append(f"total {total} tok")
    wr = j.get("result")
    if isinstance(wr, dict) and isinstance(wr.get("warnSummary"), str) and wr["warnSummary"].strip():
        lines.append("")
        lines.append(wr["warnSummary"].rstrip())
    return "\n".join(lines)


def main():
    try:
        data = json.load(sys.stdin)
    except Exception:
        sys.exit(0)
    tp = data.get("transcript_path", "")
    sid = data.get("session_id", "")
    if not tp or not sid:
        sys.exit(0)
    sess_dir = os.path.join(os.path.dirname(tp), sid)
    workflows_dir = os.path.join(sess_dir, "workflows")
    subagents_dir = os.path.join(sess_dir, "subagents")

    # State: {"watermark": <epoch_ms>, "settled": [stem,...]}. The watermark is
    # set on the FIRST Stop of the session ("started watching at T"); a run is
    # reported once its end (startTime+durationMs) >= watermark. This is
    # independent of when the journal file appears (running stub at launch vs
    # only at completion) and so reports the session's very first run, which a
    # seed-the-current-terminal-runs scheme silently absorbs. "settled" lists
    # every terminal stem already decided (reported OR excluded as pre-watermark)
    # so finished journals are never reloaded; only running journals re-parse.
    # Stamped BEFORE the workflows_dir check so the watermark lands on the first
    # Stop even when the dir/journal does not exist until a run completes.
    state_path = os.path.join(sess_dir, ".usage-reported.json")
    state = load_json(state_path)
    if not (isinstance(state, dict) and isinstance(state.get("watermark"), (int, float))):
        # First Stop (or absent/corrupt/legacy state): start watching now, emit
        # nothing. Installing the hook never dumps prior run history.
        write_state(state_path, time.time() * 1000.0, [])
        sys.exit(0)
    watermark = state["watermark"]
    settled = set(state.get("settled", []))

    if not os.path.isdir(workflows_dir):
        sys.exit(0)

    # Load only not-yet-settled journals; collect those now terminal.
    new_terminal = {}
    for p in glob.glob(os.path.join(workflows_dir, "wf_*.json")):
        stem = os.path.basename(p)[:-5]
        if stem in settled:
            continue
        j = load_json(p)
        if is_terminal(j):
            new_terminal[stem] = j
    if not new_terminal:
        sys.exit(0)

    def end_ms(j):
        st = j.get("startTime")
        return (st + j.get("durationMs", 0)) if isinstance(st, (int, float)) else None

    to_report = sorted(s for s, j in new_terminal.items() if (end_ms(j) is None or end_ms(j) >= watermark))
    settled |= set(new_terminal.keys())  # settle reported AND pre-watermark exclusions

    subagent_files = sorted(glob.glob(os.path.join(subagents_dir, "agent-*.jsonl")))
    blocks = []
    for stem in to_report:
        try:
            block = render_run(new_terminal[stem], subagent_files)
        except Exception:
            continue
        blocks.append(block)
        try:
            with open(os.path.join(sess_dir, f"run-usage-{stem}.txt"), "w") as f:
                f.write(block + "\n")
        except Exception:
            pass

    write_state(state_path, watermark, sorted(settled))

    if not blocks:
        sys.exit(0)
    msg = "\n\n".join(blocks)
    print(json.dumps({"systemMessage": msg}))
    sys.exit(0)


def write_state(path, watermark, settled):
    try:
        d = os.path.dirname(path)
        fd, tmp = tempfile.mkstemp(dir=d, prefix=".usage-reported.")
        with os.fdopen(fd, "w") as f:
            json.dump({"watermark": watermark, "settled": settled}, f)
        os.replace(tmp, path)
    except Exception:
        pass


if __name__ == "__main__":
    try:
        main()
    except Exception:
        sys.exit(0)