"""Sum the per-block token statistics from an opencode session trace.

Each conversation block in a trace ends with a line like::

     │ tokens 18 in 134 out 63628 cached 0 written $0.0000 tool-calls

This script adds up the *in*, *out* and *cached* columns across all blocks
and prints the three totals; the written-token count, the price and the
tool-call indicator are ignored.

Usage:

    python3 scripts/count_tokens.py opencode/000_opencode_session_....txt
"""

import argparse
import re
from collections.abc import Iterable
from pathlib import Path

# The stat line is preceded by a box-drawing character and ends either with
# "tool-calls" or "stop"; only the numeric columns matter.
STAT_RE = re.compile(r"tokens (\d+) in (\d+) out (\d+) cached (\d+) written")


def sum_tokens(lines: Iterable[str]) -> tuple[int, int, int]:
    """Add up the in/out/cached token counts of all stat lines.

    Args:
        lines: Iterable of trace lines.

    Returns:
        A ``(tokens_in, tokens_out, cached)`` triple of sums.
    """
    tokens_in = 0
    tokens_out = 0
    cached = 0
    for line in lines:
        match = STAT_RE.search(line)
        if not match:
            continue
        tokens_in += int(match.group(1))
        tokens_out += int(match.group(2))
        cached += int(match.group(3))
    return tokens_in, tokens_out, cached


def main() -> None:
    """Parse arguments, read the trace file and print the totals."""
    parser = argparse.ArgumentParser(
        description="Sum the token statistics from an opencode session trace."
    )
    parser.add_argument("trace_file", type=Path, help="path to the trace file")
    args = parser.parse_args()

    if not args.trace_file.is_file():
        parser.error(f"no such file: {args.trace_file}")

    with args.trace_file.open(encoding="utf-8", errors="replace") as trace:
        tokens_in, tokens_out, cached = sum_tokens(trace)

    print(f"tokens in:   {tokens_in}")
    print(f"tokens out:  {tokens_out}")
    print(f"cached:      {cached}")


if __name__ == "__main__":
    main()