codex_usage_costs.py

#!/usr/bin/env python3
"""Estimate OpenAI Codex CLI token usage and spend from local rollout files."""

from __future__ import annotations

import argparse
import json
import sys
import urllib.error
import urllib.request
from collections import defaultdict
from dataclasses import dataclass
from datetime import UTC, datetime, time, timedelta
from decimal import Decimal, InvalidOperation
from pathlib import Path
from typing import Any

import dateparser


DEFAULT_CODEX_ROOT = Path.home() / ".codex"
DEFAULT_CACHE_PATH = Path.home() / ".cache" / "codex-usage-costs" / "models.dev-api.json"
DEFAULT_PRICING_URL = "https://models.dev/api.json"
MILLION = Decimal("1000000")
LOCAL_TZINFO = datetime.now().astimezone().tzinfo or UTC
LOCAL_TZNAME = str(LOCAL_TZINFO)
ZERO = Decimal("0")
LONG_CONTEXT_THRESHOLD = 272_000
GRAPH_MODES: dict[str, dict[str, Any]] = {
    "hourly": {
        "count": 24,
        "title": "Last 24 hourly buckets",
        "unit_suffix": "h",
        "note": "Each point is one local-hour bucket. The newest bucket may be partial.",
        "selected_label": "This hour",
    },
    "daily": {
        "count": 14,
        "title": "Last 14 daily buckets",
        "unit_suffix": "day",
        "note": "Each point is one local-day bucket. Today's bucket may be partial.",
        "selected_label": "Today",
    },
    "weekly": {
        "count": 12,
        "title": "Last 12 weekly buckets",
        "unit_suffix": "week",
        "note": "Each point is one local-week bucket starting on Monday. This week's bucket may be partial.",
        "selected_label": "This week",
    },
    "monthly": {
        "count": 12,
        "title": "Last 12 monthly buckets",
        "unit_suffix": "month",
        "note": "Each point is one local-month bucket. This month's bucket may be partial.",
        "selected_label": "This month",
    },
}


@dataclass
class UsageEvent:
    file_path: str
    session_id: str | None
    timestamp: datetime | None
    response_id: str
    model: str
    input_tokens: int
    cached_input_tokens: int
    cache_write_input_tokens: int
    output_tokens: int
    reasoning_output_tokens: int
    billed_output_tokens: int
    total_tokens: int
    long_context: bool


@dataclass
class PriceBook:
    model_id: str
    input_cost_per_mtok: Decimal
    output_cost_per_mtok: Decimal
    cache_read_cost_per_mtok: Decimal
    cache_write_cost_per_mtok: Decimal
    long_input_cost_per_mtok: Decimal
    long_output_cost_per_mtok: Decimal
    long_cache_read_cost_per_mtok: Decimal
    long_cache_write_cost_per_mtok: Decimal


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Scan local Codex CLI rollout files under ~/.codex and estimate "
            "token usage and dollar cost using models.dev OpenAI pricing."
        )
    )
    parser.add_argument(
        "--root",
        type=Path,
        default=DEFAULT_CODEX_ROOT,
        help=f"Codex data root directory (default: {DEFAULT_CODEX_ROOT})",
    )
    parser.add_argument(
        "--since",
        type=str,
        help="Only include events on or after this ISO-8601 date/time or natural-language date expression. Date-only values start at 00:00:00 local time.",
    )
    parser.add_argument(
        "--until",
        type=str,
        help="Only include events before or at this ISO-8601 date/time or natural-language date expression. Date-only values include the full local day.",
    )
    parser.add_argument(
        "--pricing-url",
        default=DEFAULT_PRICING_URL,
        help=f"models.dev API URL (default: {DEFAULT_PRICING_URL})",
    )
    parser.add_argument(
        "--pricing-file",
        type=Path,
        help="Use a local models.dev api.json file instead of fetching it.",
    )
    parser.add_argument(
        "--cache-file",
        type=Path,
        default=DEFAULT_CACHE_PATH,
        help=f"Path for the cached pricing JSON (default: {DEFAULT_CACHE_PATH})",
    )
    parser.add_argument(
        "--offline",
        action="store_true",
        help="Do not fetch pricing; require --pricing-file or a cached file.",
    )
    parser.add_argument(
        "--top",
        type=int,
        default=15,
        help="How many models to show in the text report (default: 15).",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Emit machine-readable JSON instead of a text report.",
    )
    parser.add_argument(
        "--all",
        dest="all_modes",
        action="store_true",
        help="With --json, emit standard hourly/daily/weekly/monthly stats and buckets.",
    )
    args = parser.parse_args()
    if args.all_modes and not args.json:
        parser.error("--all requires --json")
    if args.all_modes and (args.since or args.until):
        parser.error("--all cannot be combined with --since/--until")
    return args


def parse_event_timestamp(value: Any) -> datetime | None:
    if not isinstance(value, str) or not value:
        return None
    text = value
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        dt = datetime.fromisoformat(text)
    except ValueError:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=UTC)
    return dt.astimezone(UTC)


def parse_when(value: str | None, *, end_of_day: bool = False) -> datetime | None:
    if value is None:
        return None
    text = value.strip()
    if not text:
        return None

    keyword_dt = parse_relative_day_keyword(text, end_of_day=end_of_day)
    if keyword_dt is not None:
        return keyword_dt

    try:
        return parse_iso_value(text, end_of_day=end_of_day)
    except ValueError:
        pass

    dt = dateparser.parse(
        text,
        settings={
            "RETURN_AS_TIMEZONE_AWARE": True,
            "TIMEZONE": LOCAL_TZNAME,
            "TO_TIMEZONE": "UTC",
        },
    )
    if dt is None:
        raise ValueError(f"could not parse {value!r} as a date/time")
    return dt.astimezone(UTC)


def parse_relative_day_keyword(text: str, *, end_of_day: bool) -> datetime | None:
    offsets = {
        "today": 0,
        "yesterday": -1,
        "tomorrow": 1,
    }
    offset = offsets.get(" ".join(text.lower().split()))
    if offset is None:
        return None
    day = (datetime.now(tz=LOCAL_TZINFO) + timedelta(days=offset)).date()
    return datetime.combine(
        day, time.max if end_of_day else time.min, tzinfo=LOCAL_TZINFO
    ).astimezone(UTC)


def parse_iso_value(text: str, *, end_of_day: bool) -> datetime:
    normalized = text
    if normalized.endswith("Z"):
        normalized = normalized[:-1] + "+00:00"
    has_time = "T" in normalized or " " in normalized
    dt = datetime.fromisoformat(normalized)
    if not has_time:
        dt = dt.replace(
            hour=23 if end_of_day else 0,
            minute=59 if end_of_day else 0,
            second=59 if end_of_day else 0,
            microsecond=999999 if end_of_day else 0,
        )
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=LOCAL_TZINFO)
    return dt.astimezone(UTC)


def start_of_bucket(dt: datetime, mode: str) -> datetime:
    if mode == "hourly":
        return dt.replace(minute=0, second=0, microsecond=0)
    if mode == "daily":
        return dt.replace(hour=0, minute=0, second=0, microsecond=0)
    if mode == "weekly":
        day_start = dt.replace(hour=0, minute=0, second=0, microsecond=0)
        return day_start - timedelta(days=day_start.weekday())
    if mode == "monthly":
        return dt.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    raise ValueError(f"unsupported graph mode {mode!r}")


def add_months(dt: datetime, months: int) -> datetime:
    month_index = (dt.month - 1) + months
    year = dt.year + month_index // 12
    month = (month_index % 12) + 1
    return dt.replace(year=year, month=month, day=1)


def shift_bucket(dt: datetime, mode: str, steps: int) -> datetime:
    if mode == "hourly":
        return dt + timedelta(hours=steps)
    if mode == "daily":
        return dt + timedelta(days=steps)
    if mode == "weekly":
        return dt + timedelta(weeks=steps)
    if mode == "monthly":
        return add_months(dt, steps)
    raise ValueError(f"unsupported graph mode {mode!r}")


def bucket_label(dt: datetime, mode: str) -> str:
    if mode == "hourly":
        return dt.strftime("%a %H:%M")
    if mode == "daily":
        return dt.strftime("%b %d")
    if mode == "weekly":
        return dt.strftime("wk %b %d")
    if mode == "monthly":
        return dt.strftime("%b %Y")
    raise ValueError(f"unsupported graph mode {mode!r}")


def bucket_range(now: datetime, mode: str) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    config = GRAPH_MODES[mode]
    count = int(config["count"])
    current_bucket_start = start_of_bucket(now, mode)
    first_bucket_start = shift_bucket(current_bucket_start, mode, -(count - 1))
    buckets: list[dict[str, Any]] = []
    for index in range(count):
        bucket_start = shift_bucket(first_bucket_start, mode, index)
        bucket_end = min(shift_bucket(bucket_start, mode, 1), now)
        buckets.append(
            {
                "start": bucket_start,
                "end": bucket_end,
                "label": bucket_label(bucket_start, mode),
            }
        )
    return buckets, config


def selected_window(now: datetime, mode: str) -> tuple[str, datetime, datetime]:
    start = start_of_bucket(now, mode)
    return str(GRAPH_MODES[mode]["selected_label"]), start, now


def all_modes_start(now: datetime) -> datetime:
    monthly_count = int(GRAPH_MODES["monthly"]["count"])
    return shift_bucket(start_of_bucket(now, "monthly"), "monthly", -(monthly_count - 1))


def read_json_file(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def fetch_pricing_json(args: argparse.Namespace) -> dict[str, Any]:
    if args.pricing_file:
        return read_json_file(args.pricing_file)

    cache_path: Path = args.cache_file
    if args.offline:
        if cache_path.exists():
            return read_json_file(cache_path)
        raise RuntimeError(
            f"offline mode requested but cache file does not exist: {cache_path}"
        )

    request = urllib.request.Request(
        args.pricing_url,
        headers={
            "User-Agent": "codex-usage-costs/1.0 (+https://models.dev)",
            "Accept": "application/json",
        },
    )
    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            payload = response.read().decode("utf-8")
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        if cache_path.exists():
            print(
                f"warning: could not fetch pricing ({exc}); using cached file {cache_path}",
                file=sys.stderr,
            )
            return read_json_file(cache_path)
        raise RuntimeError(f"failed to fetch pricing from {args.pricing_url}: {exc}") from exc

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    cache_path.write_text(payload, encoding="utf-8")
    return json.loads(payload)


def decimal_from_json(value: Any, default: Decimal) -> Decimal:
    if value is None:
        return default
    return Decimal(str(value))


def resolve_price_book(pricing_data: dict[str, Any], model: str) -> PriceBook | None:
    openai = pricing_data.get("openai")
    if not isinstance(openai, dict):
        return None
    models = openai.get("models")
    if not isinstance(models, dict):
        return None
    entry = models.get(model)
    if not isinstance(entry, dict):
        return None
    cost = entry.get("cost")
    if not isinstance(cost, dict):
        return None
    long_cost = cost.get("context_over_200k")
    if not isinstance(long_cost, dict):
        long_cost = {}

    input_cost = decimal_from_json(cost.get("input"), ZERO)
    output_cost = decimal_from_json(cost.get("output"), ZERO)
    cache_read_cost = decimal_from_json(cost.get("cache_read"), input_cost * Decimal("0.1"))
    cache_write_cost = decimal_from_json(cost.get("cache_write"), input_cost * Decimal("1.25"))
    return PriceBook(
        model_id=model,
        input_cost_per_mtok=input_cost,
        output_cost_per_mtok=output_cost,
        cache_read_cost_per_mtok=cache_read_cost,
        cache_write_cost_per_mtok=cache_write_cost,
        long_input_cost_per_mtok=decimal_from_json(long_cost.get("input"), input_cost),
        long_output_cost_per_mtok=decimal_from_json(long_cost.get("output"), output_cost),
        long_cache_read_cost_per_mtok=decimal_from_json(
            long_cost.get("cache_read"), cache_read_cost
        ),
        long_cache_write_cost_per_mtok=decimal_from_json(
            long_cost.get("cache_write"), cache_write_cost
        ),
    )


def resolve_price_book_cached(
    pricing_data: dict[str, Any],
    model: str,
    cache: dict[str, PriceBook | None],
) -> PriceBook | None:
    if model not in cache:
        cache[model] = resolve_price_book(pricing_data, model)
    return cache[model]


def intish(value: Any) -> int:
    if value in (None, ""):
        return 0
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, (int, float)):
        return int(value)
    if isinstance(value, str):
        text = value.strip()
        if not text:
            return 0
        return int(float(text))
    return 0


def extract_session_model(record: dict[str, Any]) -> str | None:
    if record.get("type") != "session_meta":
        return None
    payload = record.get("payload")
    if not isinstance(payload, dict):
        return None
    base_instructions = payload.get("base_instructions")
    if not isinstance(base_instructions, dict):
        return None
    provenance = base_instructions.get("provenance")
    if not isinstance(provenance, dict):
        return None
    model = provenance.get("model")
    if isinstance(model, str) and model:
        return model
    return None


def extract_usage_event(
    record: dict[str, Any], file_path: Path, session_id: str | None, model: str,
) -> UsageEvent | None:
    if record.get("type") != "token_usage_record":
        return None
    payload = record.get("payload")
    if not isinstance(payload, dict):
        return None
    response_id = payload.get("response_id")
    if not isinstance(response_id, str) or not response_id:
        return None
    usage = payload.get("usage")
    if not isinstance(usage, dict):
        return None

    input_tokens = intish(usage.get("input_tokens"))
    cached_input_tokens = intish(usage.get("cached_input_tokens"))
    cache_write_input_tokens = intish(usage.get("cache_write_input_tokens"))
    output_tokens = intish(usage.get("output_tokens"))
    reasoning_output_tokens = intish(usage.get("reasoning_output_tokens"))
    total_tokens = intish(usage.get("total_tokens"))
    if total_tokens <= 0:
        total_tokens = (
            input_tokens
            + cached_input_tokens
            + cache_write_input_tokens
            + output_tokens
            + reasoning_output_tokens
        )

    billed_output_tokens = output_tokens
    long_context = input_tokens > LONG_CONTEXT_THRESHOLD

    event_session_id = payload.get("session_id")
    if isinstance(event_session_id, str) and event_session_id:
        session_id = event_session_id

    return UsageEvent(
        file_path=str(file_path),
        session_id=session_id,
        timestamp=parse_event_timestamp(record.get("timestamp")),
        response_id=response_id,
        model=model,
        input_tokens=input_tokens,
        cached_input_tokens=cached_input_tokens,
        cache_write_input_tokens=cache_write_input_tokens,
        output_tokens=output_tokens,
        reasoning_output_tokens=reasoning_output_tokens,
        billed_output_tokens=billed_output_tokens,
        total_tokens=total_tokens,
        long_context=long_context,
    )


def iter_rollout_files(root: Path) -> list[Path]:
    sessions_root = root / "sessions"
    if not sessions_root.exists():
        raise RuntimeError(f"Codex sessions root does not exist: {sessions_root}")
    return sorted(sessions_root.rglob("rollout-*.jsonl"))


def load_usage_events(
    root: Path,
    since: datetime | None,
    until: datetime | None,
) -> tuple[list[UsageEvent], int, int]:
    files = iter_rollout_files(root)
    events: list[UsageEvent] = []
    parse_failures = 0
    seen_response_ids: set[str] = set()

    for file_path in files:
        session_model = "<unknown>"
        session_id: str | None = None
        with file_path.open("r", encoding="utf-8") as handle:
            for raw_line in handle:
                line = raw_line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    parse_failures += 1
                    continue

                if record.get("type") == "session_meta":
                    extracted_model = extract_session_model(record)
                    if extracted_model is not None:
                        session_model = extracted_model
                    payload = record.get("payload")
                    if isinstance(payload, dict):
                        sid = payload.get("session_id") or payload.get("id")
                        if isinstance(sid, str) and sid:
                            session_id = sid
                    continue

                event = extract_usage_event(record, file_path, session_id, session_model)
                if event is None:
                    continue
                if event.response_id in seen_response_ids:
                    continue
                seen_response_ids.add(event.response_id)
                if since and event.timestamp and event.timestamp < since:
                    continue
                if until and event.timestamp and event.timestamp > until:
                    continue
                events.append(event)

    return events, parse_failures, len(files)


def format_int(value: int) -> str:
    return f"{value:,}"


def format_decimal_usd(value: Decimal) -> str:
    return f"${value.quantize(Decimal('0.0001')):,}"


def compute_event_costs(event: UsageEvent, price_book: PriceBook) -> dict[str, Decimal]:
    long_context = event.input_tokens > LONG_CONTEXT_THRESHOLD
    input_rate = (
        price_book.long_input_cost_per_mtok if long_context else price_book.input_cost_per_mtok
    )
    output_rate = (
        price_book.long_output_cost_per_mtok if long_context else price_book.output_cost_per_mtok
    )
    cache_read_rate = (
        price_book.long_cache_read_cost_per_mtok
        if long_context
        else price_book.cache_read_cost_per_mtok
    )
    cache_write_rate = (
        price_book.long_cache_write_cost_per_mtok
        if long_context
        else price_book.cache_write_cost_per_mtok
    )
    uncached_input = max(
        0,
        event.input_tokens - event.cached_input_tokens - event.cache_write_input_tokens,
    )
    input_cost = Decimal(uncached_input) * input_rate / MILLION
    cache_read_cost = Decimal(event.cached_input_tokens) * cache_read_rate / MILLION
    cache_write_cost = Decimal(event.cache_write_input_tokens) * cache_write_rate / MILLION
    output_cost = Decimal(event.output_tokens) * output_rate / MILLION
    total_cost = input_cost + cache_read_cost + cache_write_cost + output_cost
    return {
        "input_cost_usd": input_cost,
        "cache_read_cost_usd": cache_read_cost,
        "cache_write_cost_usd": cache_write_cost,
        "output_cost_usd": output_cost,
        "total_cost_usd": total_cost,
    }


def summarize(events: list[UsageEvent], pricing_data: dict[str, Any]) -> dict[str, Any]:
    totals = {
        "responses": 0,
        "input_tokens": 0,
        "cached_input_tokens": 0,
        "cache_write_input_tokens": 0,
        "output_tokens": 0,
        "reasoning_output_tokens": 0,
        "billed_output_tokens": 0,
        "total_tokens": 0,
    }
    spend = {
        "input_cost_usd": ZERO,
        "cache_read_cost_usd": ZERO,
        "cache_write_cost_usd": ZERO,
        "output_cost_usd": ZERO,
        "total_cost_usd": ZERO,
    }
    by_model: dict[str, dict[str, Any]] = defaultdict(
        lambda: {
            "responses": 0,
            "input_tokens": 0,
            "cached_input_tokens": 0,
            "cache_write_input_tokens": 0,
            "output_tokens": 0,
            "reasoning_output_tokens": 0,
            "billed_output_tokens": 0,
            "total_tokens": 0,
            "total_cost_usd": ZERO,
            "priced_as": None,
        }
    )
    unknown_models: dict[str, int] = defaultdict(int)
    timestamps = [event.timestamp for event in events if event.timestamp is not None]
    price_book_cache: dict[str, PriceBook | None] = {}

    for event in events:
        totals["responses"] += 1
        totals["input_tokens"] += event.input_tokens
        totals["cached_input_tokens"] += event.cached_input_tokens
        totals["cache_write_input_tokens"] += event.cache_write_input_tokens
        totals["output_tokens"] += event.output_tokens
        totals["reasoning_output_tokens"] += event.reasoning_output_tokens
        totals["billed_output_tokens"] += event.billed_output_tokens
        totals["total_tokens"] += event.total_tokens

        model_row = by_model[event.model]
        model_row["responses"] += 1
        model_row["input_tokens"] += event.input_tokens
        model_row["cached_input_tokens"] += event.cached_input_tokens
        model_row["cache_write_input_tokens"] += event.cache_write_input_tokens
        model_row["output_tokens"] += event.output_tokens
        model_row["reasoning_output_tokens"] += event.reasoning_output_tokens
        model_row["billed_output_tokens"] += event.billed_output_tokens
        model_row["total_tokens"] += event.total_tokens

        price_book = resolve_price_book_cached(pricing_data, event.model, price_book_cache)
        if price_book is None:
            unknown_models[event.model] += event.total_tokens
            continue

        model_row["priced_as"] = price_book.model_id
        costs = compute_event_costs(event, price_book)
        spend["input_cost_usd"] += costs["input_cost_usd"]
        spend["cache_read_cost_usd"] += costs["cache_read_cost_usd"]
        spend["cache_write_cost_usd"] += costs["cache_write_cost_usd"]
        spend["output_cost_usd"] += costs["output_cost_usd"]
        spend["total_cost_usd"] += costs["total_cost_usd"]
        model_row["total_cost_usd"] += costs["total_cost_usd"]

    model_rows = [
        {
            "model": model,
            "priced_as": row["priced_as"],
            "responses": row["responses"],
            "input_tokens": row["input_tokens"],
            "cached_input_tokens": row["cached_input_tokens"],
            "cache_write_input_tokens": row["cache_write_input_tokens"],
            "output_tokens": row["output_tokens"],
            "reasoning_output_tokens": row["reasoning_output_tokens"],
            "billed_output_tokens": row["billed_output_tokens"],
            "total_tokens": row["total_tokens"],
            "total_cost_usd": str(row["total_cost_usd"]),
        }
        for model, row in by_model.items()
    ]
    model_rows.sort(key=lambda row: Decimal(row["total_cost_usd"]), reverse=True)

    return {
        "totals": totals,
        "spend": {key: str(value) for key, value in spend.items()},
        "by_model": model_rows,
        "unknown_models": dict(unknown_models),
        "date_range": {
            "first_event": min(timestamps).isoformat() if timestamps else None,
            "last_event": max(timestamps).isoformat() if timestamps else None,
        },
    }


def empty_mode_totals() -> dict[str, Any]:
    return {
        "responses": 0,
        "input_tokens": 0,
        "cached_input_tokens": 0,
        "cache_write_input_tokens": 0,
        "output_tokens": 0,
        "reasoning_output_tokens": 0,
        "billed_output_tokens": 0,
        "total_tokens": 0,
        "long_input_tokens": 0,
        "total_cost_usd": ZERO,
    }


def add_event_to_mode_totals(
    totals: dict[str, Any],
    event: UsageEvent,
    total_cost_usd: Decimal,
) -> None:
    totals["responses"] += 1
    totals["input_tokens"] += event.input_tokens
    totals["cached_input_tokens"] += event.cached_input_tokens
    totals["cache_write_input_tokens"] += event.cache_write_input_tokens
    totals["output_tokens"] += event.output_tokens
    totals["reasoning_output_tokens"] += event.reasoning_output_tokens
    totals["billed_output_tokens"] += event.billed_output_tokens
    totals["total_tokens"] += event.total_tokens
    if event.long_context:
        uncached = max(
            0,
            event.input_tokens - event.cached_input_tokens - event.cache_write_input_tokens,
        )
        totals["long_input_tokens"] += uncached
    totals["total_cost_usd"] += total_cost_usd


def selected_codex_token_breakdown(totals: dict[str, Any]) -> dict[str, Any]:
    cached = int(totals["cached_input_tokens"])
    cache_write = int(totals["cache_write_input_tokens"])
    long_input = int(totals["long_input_tokens"])
    short_input = max(0, int(totals["input_tokens"]) - cached - cache_write - long_input)
    return {
        "input": short_input,
        "long_input": long_input,
        "output": int(totals["output_tokens"]),
        "cache_read": cached,
        "cache_write": cache_write,
    }


def build_mode_payload(
    now: datetime,
    mode: str,
    buckets: list[dict[str, Any]],
    config: dict[str, Any],
    bucket_totals: list[dict[str, Any]],
) -> dict[str, Any]:
    selected_label, selected_start, selected_end = selected_window(now, mode)
    selected = bucket_totals[-1] if bucket_totals else empty_mode_totals()
    bucket_summaries = [
        {
            "start": bucket["start"].isoformat(),
            "end": bucket["end"].isoformat(),
            "label": bucket["label"],
            "cost_usd": str(bucket_totals[index]["total_cost_usd"]),
        }
        for index, bucket in enumerate(buckets)
    ]
    return {
        "selected": {
            "label": selected_label,
            "start": selected_start.isoformat(),
            "end": selected_end.isoformat(),
            "cost_usd": str(selected["total_cost_usd"]),
            "token_breakdown": selected_codex_token_breakdown(selected),
            "responses": int(selected["responses"]),
        },
        "graph": {
            "mode": mode,
            "title": str(config["title"]),
            "unit_suffix": str(config["unit_suffix"]),
            "note": str(config["note"]),
        },
        "buckets": bucket_summaries,
    }


def build_all_modes_payload(
    events: list[UsageEvent],
    pricing_data: dict[str, Any],
    now: datetime | None = None,
) -> dict[str, Any]:
    if now is None:
        now = datetime.now().astimezone()
    range_start = all_modes_start(now)
    relevant_events = [
        event
        for event in events
        if event.timestamp is not None and range_start <= event.timestamp <= now
    ]
    mode_state: dict[str, dict[str, Any]] = {}
    for mode in GRAPH_MODES:
        buckets, config = bucket_range(now, mode)
        mode_state[mode] = {
            "buckets": buckets,
            "config": config,
            "bucket_index": {bucket["start"]: index for index, bucket in enumerate(buckets)},
            "bucket_totals": [empty_mode_totals() for _ in buckets],
        }

    price_book_cache: dict[str, PriceBook | None] = {}
    for event in relevant_events:
        assert event.timestamp is not None
        total_cost_usd = ZERO
        price_book = resolve_price_book_cached(pricing_data, event.model, price_book_cache)
        if price_book is not None:
            total_cost_usd = compute_event_costs(event, price_book)["total_cost_usd"]

        event_local_timestamp = event.timestamp.astimezone(now.tzinfo)
        for mode, state in mode_state.items():
            bucket_start = start_of_bucket(event_local_timestamp, mode)
            bucket_index = state["bucket_index"].get(bucket_start)
            if bucket_index is None:
                continue
            add_event_to_mode_totals(
                state["bucket_totals"][bucket_index], event, total_cost_usd
            )

    return {
        "updated_at": int(now.timestamp()),
        "modes": {
            mode: build_mode_payload(
                now,
                mode,
                state["buckets"],
                state["config"],
                state["bucket_totals"],
            )
            for mode, state in mode_state.items()
        },
    }


def print_text_report(
    summary: dict[str, Any],
    root: Path,
    files_scanned: int,
    parse_failures: int,
    top: int,
) -> None:
    totals = summary["totals"]
    spend = {key: Decimal(value) for key, value in summary["spend"].items()}
    by_model = summary["by_model"]
    unknown_models = summary["unknown_models"]
    date_range = summary["date_range"]

    print(f"Codex root: {root}")
    print(f"Rollout files scanned: {files_scanned:,}")
    print(f"Codex responses counted: {totals['responses']:,}")
    print(f"Unreadable JSONL lines skipped: {parse_failures:,}")
    print("Spend is an estimate using models.dev OpenAI pricing.")
    if date_range["first_event"] and date_range["last_event"]:
        print(f"Date range: {date_range['first_event']} .. {date_range['last_event']}")
    print()

    uncached_input = max(
        0,
        totals["input_tokens"]
        - totals["cached_input_tokens"]
        - totals["cache_write_input_tokens"],
    )
    print("Tokens")
    print(f"  input (total):  {format_int(totals['input_tokens'])}")
    print(f"    uncached:     {format_int(uncached_input)}")
    print(f"    cache read:   {format_int(totals['cached_input_tokens'])}")
    print(f"    cache write:  {format_int(totals['cache_write_input_tokens'])}")
    print(f"  output (total): {format_int(totals['output_tokens'])}")
    print(f"    reasoning:    {format_int(totals['reasoning_output_tokens'])}")
    print(f"  total billed:   {format_int(totals['total_tokens'])}")
    print()

    print("Estimated spend (USD)")
    print(f"  input:          {format_decimal_usd(spend['input_cost_usd'])}")
    print(f"  cache read:     {format_decimal_usd(spend['cache_read_cost_usd'])}")
    print(f"  cache write:    {format_decimal_usd(spend['cache_write_cost_usd'])}")
    print(f"  billed output:  {format_decimal_usd(spend['output_cost_usd'])}")
    print(f"  total:          {format_decimal_usd(spend['total_cost_usd'])}")
    print()

    print("By model")
    limit = max(0, min(top, len(by_model)))
    for row in by_model[:limit]:
        print(
            "  "
            f"{row['model']}: "
            f"{format_decimal_usd(Decimal(row['total_cost_usd']))} "
            f"(responses={row['responses']:,}, "
            f"input={format_int(row['input_tokens'])}, "
            f"cache_read={format_int(row['cached_input_tokens'])}, "
            f"billed_output={format_int(row['billed_output_tokens'])})"
        )

    if unknown_models:
        print()
        print("Unpriced models")
        for model, tokens in sorted(unknown_models.items(), key=lambda item: item[1], reverse=True):
            print(f"  {model}: {format_int(tokens)} tokens")


def main() -> int:
    args = parse_args()
    pricing_data = fetch_pricing_json(args)

    if args.all_modes:
        now = datetime.now().astimezone()
        since = all_modes_start(now)
        events, parse_failures, files_scanned = load_usage_events(
            root=args.root,
            since=since,
            until=now,
        )
        payload = {
            "root": str(args.root),
            "files_scanned": files_scanned,
            "parse_failures": parse_failures,
            **build_all_modes_payload(events, pricing_data, now),
        }
        print(json.dumps(payload, indent=2, sort_keys=True))
        return 0

    since = parse_when(args.since, end_of_day=False)
    until = parse_when(args.until, end_of_day=True)
    events, parse_failures, files_scanned = load_usage_events(
        root=args.root,
        since=since,
        until=until,
    )
    summary = summarize(events, pricing_data)

    if args.json:
        print(
            json.dumps(
                {
                    "root": str(args.root),
                    "files_scanned": files_scanned,
                    "parse_failures": parse_failures,
                    **summary,
                },
                indent=2,
                sort_keys=True,
            )
        )
    else:
        print_text_report(summary, args.root, files_scanned, parse_failures, args.top)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())