#!/usr/bin/python3
# omarchy:summary=Print the Claude Code usage record as JSON
# omarchy:args=[--force] [--limits-only]
# omarchy:hidden=true
"""Collect Claude Code usage into one display-ready JSON record.

Everything the agents panel shows for Claude comes from this one
command: local transcript stats from ~/.claude/projects, the stats-cache and
history fallbacks for machines without transcripts, pi/omp and opencode
sessions that ran on an Anthropic provider, and the authoritative rate
limits from Anthropic's OAuth usage endpoint. The panel itself only ever
reads the JSON this prints; it never talks to disk formats or endpoints.
"""

from __future__ import annotations

import argparse
import datetime as dt
import fcntl
import hashlib
import json
import os
import re
import sqlite3
import sys
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any

AGENT_ID = "claude"
AGENT_NAME = "Claude Code"
AUTH_HELP = "Run `claude auth login` to restore authoritative usage."
USAGE_ENDPOINT = "https://api.anthropic.com/api/oauth/usage"
PROBE_MIN_INTERVAL_SECONDS = 15
# Anthropic rate-limits its usage endpoint readily. A refused re-check of
# numbers this recent still describes the account; only older ones are
# reported as stale.
CURRENT_ENOUGH_SECONDS = 900


def config_dir() -> Path:
  return expand_path(os.environ.get("CLAUDE_CONFIG_DIR") or "~/.claude")


def expand_path(value: str) -> Path:
  return Path(os.path.expandvars(os.path.expanduser(value))).resolve()


def cache_root() -> Path:
  root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage"
  root.mkdir(parents=True, exist_ok=True)
  return root


def date_string(value: dt.date) -> str:
  return value.strftime("%Y-%m-%d")


def recent_date_strings() -> list[str]:
  today = dt.datetime.now().date()
  return [date_string(today - dt.timedelta(days=offset)) for offset in range(6, -1, -1)]


def local_date_string() -> str:
  return date_string(dt.datetime.now().date())


def local_date_from_timestamp(value: Any) -> str:
  if value is None:
    return local_date_string()

  if isinstance(value, (int, float)):
    try:
      seconds = float(value) / 1000.0 if float(value) > 10_000_000_000 else float(value)
      return date_string(dt.datetime.fromtimestamp(seconds).date())
    except Exception:
      return local_date_string()

  raw = str(value).strip()
  if not raw:
    return local_date_string()

  # Claude JSONL timestamps are usually ISO-8601. Python accepts offsets but
  # not a trailing Z until we normalize it to +00:00.
  try:
    parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
    if parsed.tzinfo is not None:
      parsed = parsed.astimezone()
    return date_string(parsed.date())
  except Exception:
    return local_date_string()


def usage_token(usage: dict[str, Any], snake_key: str, camel_key: str) -> int:
  value = usage.get(snake_key, usage.get(camel_key, 0))
  try:
    return round(float(value or 0))
  except Exception:
    return 0


def number(value: Any) -> int:
  try:
    n = float(value or 0)
    return round(n) if n == n else 0
  except Exception:
    return 0


def empty_bucket() -> dict[str, int]:
  return {
    "inputTokens": 0,
    "outputTokens": 0,
    "cacheReadInputTokens": 0,
    "cacheCreationInputTokens": 0,
  }


# ---------------------------------------------------------------- local scan


# Transcripts are append-only, and most of the corpus never changes between
# ticks, so the scan keeps a per-file index: the unique usage records already
# parsed out of each file plus the byte offset they end at. A file whose size
# and mtime match the index is not opened; a file that grew is read from the
# stored offset; anything else (truncated, rewritten, replaced, unknown) is
# read from the start. Without this the whole history was re-parsed every
# refresh, and the cost grew for the life of the machine.

INDEX_VERSION = 4

# How much of each end of what was read the index remembers a hash of, so a
# file replaced or rewritten with more than it had is not mistaken for an
# append.
INDEX_HEAD_BYTES = 4096


# Records hold local days, so an index is only true in the timezone that
# wrote it; after a change every transcript is read again.
def index_zone() -> list[Any]:
  return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone]


def load_index(index_file: Path) -> dict[str, Any]:
  try:
    data = json.loads(index_file.read_text(encoding="utf-8"))
    if (
      isinstance(data, dict)
      and data.get("version") == INDEX_VERSION
      and data.get("zone") == index_zone()
      and isinstance(data.get("files"), dict)
    ):
      return data["files"]
  except Exception:
    pass
  return {}


# What was already read, by its first and last few kilobytes: a transcript
# rewritten in place rather than appended to changes one or the other.
def read_digest(path: Path, length: int) -> str:
  digest = hashlib.sha1()
  with path.open("rb") as handle:
    digest.update(handle.read(min(length, INDEX_HEAD_BYTES)))
    handle.seek(max(0, length - INDEX_HEAD_BYTES))
    digest.update(handle.read(min(length, INDEX_HEAD_BYTES)))
  return digest.hexdigest()


def parse_usage_records(path: Path, entry: dict[str, Any] | None) -> dict[str, Any]:
  """Return the index entry for one transcript, reusing `entry` where the file
  only grew. Records are [key, model, day, input, output, cacheRead,
  cacheWrite, session]; `day` is None when the line carries no timestamp."""
  stat = path.stat()
  size = stat.st_size
  mtime = stat.st_mtime_ns

  if entry and entry.get("size") == size and entry.get("mtime") == mtime:
    return entry

  # Only the same file, grown past what was read with what was read still in
  # place, is an append. A transcript replaced by a larger one is a new file,
  # and one rewritten in place has a different start or end.
  if (
    entry
    and size > number(entry.get("size"))
    and entry.get("file") == [stat.st_dev, stat.st_ino]
    and entry.get("head") == read_digest(path, number(entry.get("offset")))
  ):
    records = list(entry.get("records") or [])
    offset = number(entry.get("offset"))
    line_number = number(entry.get("lines"))
  else:
    records = []
    offset = 0
    line_number = 0

  # A streamed response is several lines sharing one message id, and the
  # first line's output_tokens is a placeholder while the last carries the
  # real count. Keep each message's line with the highest output count (the
  # last one read wins a tie), whole: a response that fell back to another
  # model mid-stream has lines that are separate snapshots, and mixing fields
  # across them over-counts cache tokens and credits the wrong model.
  positions = {record[0]: index for index, record in enumerate(records)}
  with path.open("rb") as handle:
    handle.seek(offset)
    for raw in handle:
      if not raw.endswith(b"\n"):
        # A line still being written; pick it up whole next time.
        break
      offset += len(raw)
      line_number += 1

      # Cheap pre-filter before JSON parsing keeps files with unrelated
      # lines inexpensive.
      if b'"usage":' not in raw:
        continue

      try:
        parsed = json.loads(raw.decode("utf-8", errors="replace"))
      except Exception:
        continue

      message = parsed.get("message") if isinstance(parsed.get("message"), dict) else {}
      if parsed.get("type") != "assistant" and message.get("role") != "assistant":
        continue

      usage = message.get("usage") or parsed.get("usage")
      if not isinstance(usage, dict):
        continue

      message_id = message.get("id") or parsed.get("messageId") or ""
      key = str(message_id) if message_id else f"{path}:{parsed.get('uuid') or parsed.get('requestId') or line_number}"

      input_tokens = usage_token(usage, "input_tokens", "inputTokens")
      output_tokens = usage_token(usage, "output_tokens", "outputTokens")
      cache_read = usage_token(usage, "cache_read_input_tokens", "cacheReadInputTokens")
      cache_write = usage_token(usage, "cache_creation_input_tokens", "cacheCreationInputTokens")
      if input_tokens + output_tokens + cache_read + cache_write <= 0:
        continue

      held = positions.get(key)
      if held is not None and output_tokens < records[held][4]:
        continue

      timestamp = parsed.get("timestamp") or message.get("timestamp")
      record = [
        key,
        str(message.get("model") or parsed.get("model") or "claude"),
        local_date_from_timestamp(timestamp) if timestamp is not None else None,
        input_tokens,
        output_tokens,
        cache_read,
        cache_write,
        str(parsed.get("sessionId") or path),
      ]
      if held is None:
        positions[key] = len(records)
        records.append(record)
      else:
        records[held] = record

  return {
    "size": size,
    "file": [stat.st_dev, stat.st_ino],
    "mtime": mtime,
    "head": read_digest(path, offset),
    "offset": offset,
    "lines": line_number,
    "records": records,
  }


def scan_projects(projects_path: Path, index_file: Path | None = None) -> dict[str, Any]:
  previous = load_index(index_file) if index_file else {}
  current: dict[str, Any] = {}

  # Directory order, as before: when one API message shows up in several
  # transcripts with the same output count, the last file visited wins, and
  # the numbers must not move because the scan got faster.
  order: list[str] = []
  unreadable_count = 0
  unreadable_examples: list[str] = []
  files = projects_path.rglob("*.jsonl") if projects_path.is_dir() else []
  for path in files:
    key = str(path)
    try:
      current[key] = parse_usage_records(path, previous.get(key))
    except Exception as exc:
      unreadable_count += 1
      if len(unreadable_examples) < 3:
        unreadable_examples.append(f"{path}: {exc}")
      # What was read before is still true of the file.
      if key in previous:
        current[key] = previous[key]
    if key in current:
      order.append(key)

  if unreadable_count:
    examples = "; ".join(unreadable_examples)
    print(
      f"Ignoring {unreadable_count} unreadable Claude project files"
      f" (examples: {examples})",
      file=sys.stderr,
    )

  if index_file and current != previous:
    try:
      write_json(index_file, {"version": INDEX_VERSION, "zone": index_zone(), "files": current})
    except Exception as exc:
      print(f"Could not write Claude scan index {index_file}: {exc}", file=sys.stderr)

  return summarize_records(current, order)


def summarize_records(indexed: dict[str, Any], order: list[str]) -> dict[str, Any]:
  today = local_date_string()
  recent_dates = recent_date_strings()
  recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}

  # One API message can show up in several transcripts (a resumed session
  # carries the earlier ones). As within a file, the record with the highest
  # output count wins, the last one visited breaking a tie.
  seen: dict[str, list[Any]] = {}
  sessions: set[str] = set()
  active_days: set[str] = set()
  today_sessions: set[str] = set()
  today_tokens: dict[str, int] = {}
  usage_by_model: dict[str, dict[str, int]] = {}
  prompts = 0
  today_prompt_count = 0
  today_token_total = 0

  for key in order:
    for record in indexed[key].get("records") or []:
      held = seen.get(record[0])
      if held is not None and record[4] < held[4]:
        continue
      seen[record[0]] = record

  for _, model, day, input_tokens, output_tokens, cache_read, cache_write, session_key in seen.values():
    if day is None:
      day = today
    total = input_tokens + output_tokens + cache_read + cache_write

    sessions.add(session_key)
    active_days.add(day)
    prompts += 1

    bucket = usage_by_model.setdefault(model, empty_bucket())
    bucket["inputTokens"] += input_tokens
    bucket["outputTokens"] += output_tokens
    bucket["cacheReadInputTokens"] += cache_read
    bucket["cacheCreationInputTokens"] += cache_write

    if day in recent:
      # recentDays.messageCount is actually a token total, despite the
      # legacy name shared with synced snapshots.
      recent[day]["messageCount"] += total

    if day == today:
      today_prompt_count += 1
      today_sessions.add(session_key)
      today_token_total += total
      today_tokens[model] = today_tokens.get(model, 0) + total

  return {
    "todayPrompts": today_prompt_count,
    "todaySessions": len(today_sessions),
    "todayTotalTokens": today_token_total,
    "todayTokensByModel": today_tokens,
    "recentDays": [recent[day] for day in recent_dates],
    "modelUsage": usage_by_model,
    "totalPrompts": prompts,
    "totalSessions": len(sessions),
    # Days with any recorded usage, for the all-time "N days" summary. The
    # dates travel too: merging snapshots from several machines needs their
    # union, which a count alone cannot give.
    "activeDays": len(active_days),
    "activeDates": sorted(active_days),
  }


def scan_cache_paths(projects_path: Path) -> tuple[Path, Path, Path]:
  digest = hashlib.sha1(str(projects_path).encode("utf-8")).hexdigest()[:16]
  root = cache_root()
  return root / f"claude-scan-{digest}.json", root / f"claude-scan-{digest}.lock", root / f"claude-index-{digest}.json"


def read_json(path: Path) -> dict[str, Any] | None:
  try:
    return json.loads(path.read_text(encoding="utf-8"))
  except Exception:
    return None


def read_fresh_json(path: Path, max_age_seconds: float) -> dict[str, Any] | None:
  if max_age_seconds <= 0 or not path.exists():
    return None
  try:
    # A negative age means the mtime is in the future: the clock moved
    # backwards since the write, so the cache's freshness cannot be trusted.
    age = time.time() - path.stat().st_mtime
    if 0 <= age <= max_age_seconds:
      return read_json(path)
  except Exception:
    return None
  return None


def write_json(path: Path, payload: dict[str, Any]) -> None:
  # A temp name unique to this writer, not derived from the target: several
  # collectors can run at once (the update command backgrounds one per agent,
  # the panel refreshes on its own), and a shared temp path means the second
  # replace finds the first one's file already moved away.
  handle_fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp")
  tmp = Path(tmp_name)
  try:
    with os.fdopen(handle_fd, "w", encoding="utf-8") as handle:
      handle.write(json.dumps(payload, separators=(",", ":"), sort_keys=True) + "\n")
    # mkstemp opens at 0600; these caches were world-readable before.
    tmp.chmod(0o644)
    tmp.replace(path)
  except BaseException:
    tmp.unlink(missing_ok=True)
    raise


def cached_scan(projects_path: Path, max_age_seconds: float, use_index: bool = True) -> dict[str, Any]:
  cache_file, lock_file, index_file = scan_cache_paths(projects_path)

  cached = read_fresh_json(cache_file, max_age_seconds)
  if cached is not None:
    return cached

  with lock_file.open("w") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    cached = read_fresh_json(cache_file, max_age_seconds)
    if cached is not None:
      return cached
    if not use_index:
      index_file.unlink(missing_ok=True)
    summary = scan_projects(projects_path, index_file)
    write_json(cache_file, summary)
    return summary


# ------------------------------------------------------------- local fallback
#
# A machine without transcripts on disk can still know its history: Claude
# Code keeps aggregate counters in stats-cache.json and per-prompt history in
# history.jsonl. Only consulted when the project scan comes back empty.


def stats_cache_fallback(claude_dir: Path) -> dict[str, Any] | None:
  try:
    data = json.loads((claude_dir / "stats-cache.json").read_text(encoding="utf-8"))
  except Exception:
    return None

  today = local_date_string()
  daily_model_tokens = data.get("dailyModelTokens") or []
  today_tokens = {}
  for entry in daily_model_tokens:
    if isinstance(entry, dict) and entry.get("date") == today:
      today_tokens = entry.get("tokensByModel") or {}
      break

  daily_activity = [day for day in (data.get("dailyActivity") or []) if isinstance(day, dict)]
  active_dates = sorted({str(day.get("date")) for day in daily_activity if number(day.get("messageCount")) > 0 and day.get("date")})
  today_prompts, today_sessions = today_prompts_from_history(claude_dir)

  return {
    "todayPrompts": today_prompts,
    "todaySessions": today_sessions,
    "todayTotalTokens": sum(number(v) for v in today_tokens.values()),
    "todayTokensByModel": today_tokens,
    "recentDays": daily_activity[-7:],
    "modelUsage": data.get("modelUsage") or {},
    "totalPrompts": number(data.get("totalMessages")),
    "totalSessions": number(data.get("totalSessions")),
    "activeDays": len(active_dates),
    "activeDates": active_dates,
  }


def today_prompts_from_history(claude_dir: Path) -> tuple[int, int]:
  prompts = 0
  sessions: set[str] = set()
  start_of_day = dt.datetime.combine(dt.datetime.now().date(), dt.time.min).timestamp() * 1000
  try:
    with (claude_dir / "history.jsonl").open("r", encoding="utf-8", errors="replace") as handle:
      lines = handle.readlines()
  except Exception:
    return 0, 0

  for line in reversed(lines):
    line = line.strip()
    if not line:
      continue
    try:
      entry = json.loads(line)
    except Exception:
      continue
    if number(entry.get("timestamp")) < start_of_day:
      break
    prompts += 1
    if entry.get("sessionId"):
      sessions.add(str(entry.get("sessionId")))
  return prompts, len(sessions)


# --------------------------------------------------------------- pi and omp
#
# These agents can consume a Claude subscription without writing native
# Claude Code transcripts. Their compatible JSONL session formats carry the
# provider, model, and token usage on every assistant message.


# `omp --profile=<name>` (and pi's equivalent) relocates the whole agent tree
# under <base>/profiles/<name>/, so a subscription driven entirely through a
# profile leaves the default root empty and its usage uncounted. Sessions are
# keyed by file path, so a profile adds sessions rather than double-counting
# the default root.
def pi_session_roots() -> list[Path]:
  roots: list[Path] = []
  for base in (Path.home() / ".pi", Path.home() / ".omp"):
    roots.append(base / "agent" / "sessions")
    profiles = base / "profiles"
    try:
      # Sorted so the scan order does not depend on directory order.
      roots.extend(sorted(child / "agent" / "sessions" for child in profiles.iterdir() if child.is_dir()))
    except OSError:
      # No profiles directory, or it is unreadable: the default root stands.
      pass
  return roots


def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None:
  roots = pi_session_roots()
  cache_file = cache_root() / "claude-pi-sessions.json"
  cached = read_fresh_json(cache_file, max_age_seconds)
  if cached is not None:
    return cached.get("stats")

  today = local_date_string()
  recent_dates = recent_date_strings()
  recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
  sessions: set[str] = set()
  active_days: set[str] = set()
  today_sessions: set[str] = set()
  today_tokens: dict[str, int] = {}
  usage_by_model: dict[str, dict[str, int]] = {}
  seen: set[str] = set()
  prompts = 0
  today_prompt_count = 0
  today_token_total = 0

  for root in roots:
    files = root.rglob("*.jsonl") if root.is_dir() else []
    for path in files:
      try:
        with path.open("r", encoding="utf-8", errors="replace") as handle:
          for line_number, line in enumerate(handle, 1):
            if '"usage"' not in line or '"assistant"' not in line:
              continue
            try:
              entry = json.loads(line)
              message = entry.get("message") if isinstance(entry.get("message"), dict) else {}
              if entry.get("type") != "message" or message.get("role") != "assistant":
                continue
              provider = str(message.get("provider") or "")
              if provider != "anthropic":
                continue
              unique_key = f"{path}:{entry.get('id') or line_number}"
              if unique_key in seen:
                continue
              seen.add(unique_key)
              usage = message.get("usage") or {}
              input_tokens = usage_token(usage, "input", "inputTokens")
              output_tokens = usage_token(usage, "output", "outputTokens")
              cache_read = usage_token(usage, "cacheRead", "cache_read_input_tokens")
              cache_write = usage_token(usage, "cacheWrite", "cache_creation_input_tokens")
              total = input_tokens + output_tokens + cache_read + cache_write
              if total <= 0:
                total = number(usage.get("totalTokens"))
                input_tokens = total
              if total <= 0:
                continue
              model = str(message.get("model") or "claude")
              day = local_date_from_timestamp(entry.get("timestamp") or message.get("timestamp"))
            except Exception:
              continue

            session_key = str(path)
            sessions.add(session_key)
            active_days.add(day)
            prompts += 1
            bucket = usage_by_model.setdefault(model, empty_bucket())
            bucket["inputTokens"] += input_tokens
            bucket["outputTokens"] += output_tokens
            bucket["cacheReadInputTokens"] += cache_read
            bucket["cacheCreationInputTokens"] += cache_write
            if day in recent:
              recent[day]["messageCount"] += total
            if day == today:
              today_prompt_count += 1
              today_sessions.add(session_key)
              today_token_total += total
              today_tokens[model] = today_tokens.get(model, 0) + total
      except OSError:
        continue

  stats = None
  if prompts > 0:
    stats = {
      "todayPrompts": today_prompt_count,
      "todaySessions": len(today_sessions),
      "todayTotalTokens": today_token_total,
      "todayTokensByModel": today_tokens,
      "recentDays": [recent[day] for day in recent_dates],
      "modelUsage": usage_by_model,
      "totalPrompts": prompts,
      "totalSessions": len(sessions),
      "activeDays": len(active_days),
      "activeDates": sorted(active_days),
    }
  write_json(cache_file, {"stats": stats})
  return stats


# ---------------------------------------------------------------- opencode
#
# A Claude subscription burned entirely through opencode never writes a
# transcript under ~/.claude, but opencode records per-message provider,
# model, and token usage in its own database. Scan it for Anthropic-provider
# messages and merge the result into whatever the transcript scan found.


def scan_opencode_usage(max_age_seconds: float) -> dict[str, Any] | None:
  db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db"
  if not db.is_file():
    return None

  # Same freshness contract as the transcript scan: --limits-only promises to
  # reuse recent local stats, and a big opencode history walked on every panel
  # open would break that promise.
  cache_file = cache_root() / f"claude-opencode-{hashlib.sha1(str(db).encode('utf-8')).hexdigest()[:16]}.json"
  cached = read_fresh_json(cache_file, max_age_seconds)
  if cached is not None:
    return cached.get("stats")

  today = local_date_string()
  recent_dates = recent_date_strings()
  recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
  sessions: set[str] = set()
  active_days: set[str] = set()
  today_sessions: set[str] = set()
  today_tokens: dict[str, int] = {}
  usage_by_model: dict[str, dict[str, int]] = {}
  prompts = 0
  today_prompt_count = 0
  today_token_total = 0

  try:
    # Read-only: opencode may be writing right now.
    conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2)
  except sqlite3.Error:
    return None
  try:
    conn.execute("PRAGMA query_only = ON")
    tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")}
    # opencode v1 keeps messages in `message`, with the role, provider and model
    # in each row's data. v2 writes `session_message` instead: the role moves to
    # the `type` column and the provider and model nest under data.model. A
    # database can hold either or both.
    queries = []
    if "message" in tables:
      queries.append(("opencode:", "SELECT session_id, data, NULL FROM message"))
    if "session_message" in tables:
      queries.append(("opencode-v2:", "SELECT session_id, data, type FROM session_message"))
    rows = ((prefix, row) for prefix, query in queries for row in conn.execute(query))
    for session_prefix, (session_id, raw, row_type) in rows:
      # One malformed row must not abort the scan, so every shape assumption
      # lives inside the try.
      try:
        entry = json.loads(raw)
        if not isinstance(entry, dict):
          continue
        if row_type is None:
          role = entry.get("role")
          provider = entry.get("providerID")
          model_id = entry.get("modelID")
        else:
          role = row_type
          model_info = entry.get("model") or {}
          provider = model_info.get("providerID")
          model_id = model_info.get("id")
        # Exact match: opencode provider ids are free-form, and a custom
        # "anthropic-proxy" gateway is not this subscription.
        if role != "assistant" or str(provider or "") != "anthropic":
          continue
        tokens = entry.get("tokens") or {}
        cache = tokens.get("cache") or {}
        input_tokens = number(tokens.get("input"))
        # opencode keeps thinking tokens out of output; both are generated.
        output_tokens = number(tokens.get("output")) + number(tokens.get("reasoning"))
        cache_read = number(cache.get("read"))
        cache_write = number(cache.get("write"))
        total = input_tokens + output_tokens + cache_read + cache_write
        if total <= 0:
          continue

        created = number((entry.get("time") or {}).get("created"))
        day = dt.datetime.fromtimestamp(created / 1000).strftime("%Y-%m-%d") if created > 0 else today
        model = str(model_id or "claude").rstrip("/").split("/")[-1]
      except Exception:
        continue
      session_key = session_prefix + str(session_id)
      sessions.add(session_key)
      active_days.add(day)
      prompts += 1

      bucket = usage_by_model.setdefault(model, empty_bucket())
      bucket["inputTokens"] += input_tokens
      bucket["outputTokens"] += output_tokens
      bucket["cacheReadInputTokens"] += cache_read
      bucket["cacheCreationInputTokens"] += cache_write

      if day in recent:
        recent[day]["messageCount"] += total
      if day == today:
        today_prompt_count += 1
        today_sessions.add(session_key)
        today_token_total += total
        today_tokens[model] = today_tokens.get(model, 0) + total
  except sqlite3.Error:
    return None
  finally:
    conn.close()

  stats = None
  if prompts > 0:
    stats = {
      "todayPrompts": today_prompt_count,
      "todaySessions": len(today_sessions),
      "todayTotalTokens": today_token_total,
      "todayTokensByModel": today_tokens,
      "recentDays": [recent[day] for day in recent_dates],
      "modelUsage": usage_by_model,
      "totalPrompts": prompts,
      "totalSessions": len(sessions),
      "activeDays": len(active_days),
      "activeDates": sorted(active_days),
    }
  write_json(cache_file, {"stats": stats})
  return stats


def merge_stats(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
  merged = dict(base)
  for key in ("todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions"):
    merged[key] = number(base.get(key)) + number(extra.get(key))

  combined = dict(base.get("todayTokensByModel") or {})
  for model, count in (extra.get("todayTokensByModel") or {}).items():
    combined[model] = number(combined.get(model)) + number(count)
  merged["todayTokensByModel"] = combined

  usage = {model: dict(bucket) for model, bucket in (base.get("modelUsage") or {}).items()}
  for model, bucket in (extra.get("modelUsage") or {}).items():
    target = usage.setdefault(model, empty_bucket())
    for field, count in (bucket or {}).items():
      target[field] = number(target.get(field)) + number(count)
  merged["modelUsage"] = usage

  by_date: dict[str, int] = {}
  for source in (base.get("recentDays") or [], extra.get("recentDays") or []):
    for day in source:
      date = str((day or {}).get("date") or "")
      if date:
        by_date[date] = by_date.get(date, 0) + number((day or {}).get("messageCount"))
  merged["recentDays"] = [{"date": date, "messageCount": by_date[date]} for date in sorted(by_date)]

  # Sources overlap in time, so union dates rather than summing counts. A
  # fallback that only knows a count still bounds the answer from below.
  dates = set(base.get("activeDates") or []) | set(extra.get("activeDates") or [])
  merged["activeDates"] = sorted(dates)
  merged["activeDays"] = max(len(dates), number(base.get("activeDays")), number(extra.get("activeDays")))
  return merged


# ------------------------------------------------------------------- limits


# The access token, its expiry, and the display-safe plan label from the
# CLI's login. Nothing else leaves the credential store: the token goes
# nowhere but the Authorization header of the limits probe, and only the
# plan label may travel into the printed record.
def oauth_login(claude_dir: Path) -> tuple[str, int, int, str]:
  try:
    data = json.loads((claude_dir / ".credentials.json").read_text(encoding="utf-8"))
  except Exception:
    return "", 0, 0, ""
  login = data.get("claudeAiOauth")
  if not isinstance(login, dict):
    return "", 0, 0, ""
  # The profile's tier is asked first, but only one that names a multiplier
  # outranks the credentials'. The subscription still leads the label, so a
  # Team seat keeps reading Team.
  tier = profile_tier(claude_dir)
  if not re.search(r"max_(\d+x)", tier, re.IGNORECASE):
    tier = str(login.get("rateLimitTier") or "")
  plan = plan_label(tier, str(login.get("subscriptionType") or ""))
  return (
    str(login.get("accessToken") or ""),
    number(login.get("expiresAt")),
    number(login.get("refreshTokenExpiresAt")),
    plan,
  )


# The credentials carry the tier their token was minted with, and an upgrade
# does not restamp it: a refresh rewrites the file and keeps the old value, so
# an account moved to Max 20x keeps reading Max 5x. The CLI's own profile is
# refreshed against the account, so the tier it states there is asked first.
def profile_tier(claude_dir: Path) -> str:
  try:
    account = json.loads(profile_path(claude_dir).read_text(encoding="utf-8")).get("oauthAccount")
  except Exception:
    return ""
  if not isinstance(account, dict):
    return ""
  # A seat with no tier of its own is limited by its organization, which is
  # the shape a personal Max subscription takes.
  for key in ("userRateLimitTier", "organizationRateLimitTier"):
    tier = str(account.get(key) or "")
    if tier:
      return tier
  return ""


# The CLI's profile for a config home: a legacy .config.json inside it, else
# ~/.claude.json for the default home and <home>/.claude.json for one it was
# pointed at with CLAUDE_CONFIG_DIR. A copy anywhere else is never refreshed,
# and another home's profile describes another account.
def profile_path(home: Path) -> Path:
  legacy = home / ".config.json"
  if legacy.exists():
    return legacy
  configured = os.environ.get("CLAUDE_CONFIG_DIR")
  if home == expand_path("~/.claude") and not (configured and expand_path(configured) == home):
    return Path.home() / ".claude.json"
  return home / ".claude.json"


# The rate-limit tier names the ceiling, not the plan: a Team premium seat
# runs on `default_claude_max_5x` while paying for Team. Lead with the
# subscription and keep the multiplier as its qualifier, so Max still reads
# "Max 5x" and a Team seat reads "Team 5x" rather than passing as Max.
def plan_label(tier: str, subscription: str) -> str:
  plan = subscription[0].upper() + subscription[1:] if subscription else ""
  match = re.search(r"max_(\d+x)", tier, re.IGNORECASE) if tier else None
  if not match:
    return plan
  return (plan or "Max") + " " + match.group(1)


def parse_utilization(value: Any) -> float:
  try:
    return float(str(value).strip().replace("%", ""))
  except Exception:
    return float("nan")


def normalize_utilization(value: Any, percent_scale: bool) -> float:
  n = parse_utilization(value)
  if not (n >= 0):
    return -1.0
  # Anthropic's OAuth usage endpoint currently reports percentages (for
  # example 37.0 or 1.0). Older payloads sometimes used fractions (0.37).
  # A payload containing any value >= 1 is percent-scaled, so 1.0 renders
  # as 1%, not 100%.
  if percent_scale or n > 1:
    return min(1.0, n / 100.0)
  return min(1.0, n)


def normalize_reset_at(value: Any) -> str:
  if value is None:
    return ""
  raw = str(value).strip()
  if raw == "":
    return ""
  if raw.isdigit():
    ts = int(raw)
    if ts < 1e12:
      ts *= 1000
    try:
      return dt.datetime.fromtimestamp(ts / 1000, dt.timezone.utc).isoformat()
    except Exception:
      return raw
  try:
    parsed = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
    return parsed.isoformat()
  except Exception:
    return raw


def usage_bucket(payload: dict[str, Any], key: str) -> dict[str, Any] | None:
  bucket = payload.get(key)
  return bucket if isinstance(bucket, dict) else None


# An entry's `kind` names its window the way the flat buckets' keys do
# ("weekly_scoped", "five_hour_scoped"). The panel reads a window out of free
# text, which cannot survive a model name like "Opus 5 (1M context)" — the
# "1M" reads as a one-minute window — so the window is settled here instead
# and travels as an explicit title. It is capitalized the way the flat windows
# title themselves, so "Fable Weekly" sits beside "Weekly" rather than under it.
def scoped_window(kind: str) -> str:
  text = kind.lower()
  if "month" in text:
    return "Monthly"
  if "week" in text or "day" in text:
    return "Weekly"
  if "hour" in text or "session" in text:
    return "Session"
  return ""


# Alongside the flat buckets, the payload carries a `limits` array, and that
# array is the only place a model-scoped allowance shows up — a weekly window
# that only Fable draws from, say. The matching legacy keys
# (`seven_day_opus`, `seven_day_sonnet`, …) stayed behind at null, so a
# collector that reads buckets alone silently drops a limit the account is
# actually spending against. A model can hold more than one scoped window, and
# only the pair of model and window tells them apart, so both make the title
# and both make the key that keeps a repeat out.
def scoped_limits(payload: dict[str, Any], percent_scale: bool) -> list[dict[str, Any]]:
  entries = payload.get("limits")
  if not isinstance(entries, list):
    return []
  out: list[dict[str, Any]] = []
  seen: set[tuple[str, str]] = set()
  for entry in entries:
    if not isinstance(entry, dict):
      continue
    scope = entry.get("scope")
    model = scope.get("model") if isinstance(scope, dict) else None
    if not isinstance(model, dict):
      continue
    # A display name is what the panel wants, but an entry carrying only an id
    # still names a window worth showing.
    name = str(model.get("display_name") or model.get("id") or "").strip()
    kind = str(entry.get("kind") or "").strip()
    if name == "" or (name, kind) in seen:
      continue
    percent = normalize_utilization(entry.get("percent"), percent_scale)
    if percent < 0:
      continue
    seen.add((name, kind))
    window = scoped_window(kind)
    title = name + " " + window if window else name
    out.append({
      "label": title,
      "title": title,
      "percent": percent,
      "resetsAt": normalize_reset_at(entry.get("resets_at")),
    })
  return out


def probe_limits(access_token: str) -> dict[str, Any]:
  request = urllib.request.Request(
    USAGE_ENDPOINT,
    headers={
      "Authorization": "Bearer " + access_token,
      "anthropic-beta": "oauth-2025-04-20",
      "Accept": "application/json",
    },
  )
  try:
    with urllib.request.urlopen(request, timeout=10) as response:
      payload = json.loads(response.read().decode("utf-8", errors="replace"))
  except urllib.error.HTTPError as error:
    retry_after = error.headers.get("retry-after", "") if error.headers else ""
    if error.code == 429:
      help_text = "Anthropic's usage endpoint is rate limiting checks right now" + (
        f" (retry after {retry_after}s)" if retry_after else ""
      ) + ". Local Claude Code stats are still shown."
    else:
      help_text = f"Anthropic's usage endpoint returned status {error.code}. Local Claude Code stats are still shown."
    return {"ok": False, "helpText": help_text}
  except Exception:
    # A transport failure reached no server at all — no route, no DNS. Any
    # real answer, including an error status, is a server we should stop
    # pestering; this is not.
    return {
      "ok": False,
      "transport": True,
      "helpText": "Couldn't reach Anthropic's usage endpoint. Retrying shortly. Local Claude Code stats are still shown.",
    }

  weekly = usage_bucket(payload, "seven_day_oauth_apps") or usage_bucket(payload, "seven_day")
  session = usage_bucket(payload, "five_hour")
  raw = [session.get("utilization") if session else None, weekly.get("utilization") if weekly else None]
  # One payload speaks one convention, so the scoped entries settle the scale
  # alongside the buckets rather than assuming their own.
  entries = payload.get("limits")
  if isinstance(entries, list):
    raw += [entry.get("percent") for entry in entries if isinstance(entry, dict)]
  percent_scale = any(parse_utilization(v) >= 1 for v in raw)

  limits = []
  if session is not None:
    percent = normalize_utilization(session.get("utilization"), percent_scale)
    if percent >= 0:
      limits.append({"label": "Session (5-hour)", "percent": percent, "resetsAt": normalize_reset_at(session.get("resets_at"))})
  if weekly is not None:
    percent = normalize_utilization(weekly.get("utilization"), percent_scale)
    if percent >= 0:
      limits.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": normalize_reset_at(weekly.get("resets_at"))})
  limits.extend(scoped_limits(payload, percent_scale))

  if not limits:
    return {"ok": False, "helpText": "Anthropic's usage endpoint returned no limits. Local Claude Code stats are still shown."}
  return {"ok": True, "limits": limits}


# Cached windows reset to zero when their measured period ends. Without a
# fresh probe, the next rolling window has no known reset time yet.
def limit_window_open(entry: dict[str, Any], now: dt.datetime) -> bool:
  raw = str(entry.get("resetsAt") or "")
  if raw == "":
    return True
  try:
    resets_at = dt.datetime.fromisoformat(raw.replace("Z", "+00:00"))
  except Exception:
    return True
  if resets_at.tzinfo is None:
    resets_at = resets_at.replace(tzinfo=dt.timezone.utc)
  return resets_at > now


def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]:
  entries = cached.get("limits")
  if not isinstance(entries, list):
    return []
  now = dt.datetime.now(dt.timezone.utc)
  return [
    entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="")
    for entry in entries if isinstance(entry, dict)
  ]


def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: int, force: bool, cache_name: str = "claude-limits.json") -> dict[str, Any]:
  # `live` says the numbers came from Anthropic just now (or within the reuse
  # window) rather than from a cache kept past a failed or impossible probe.
  result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False}

  # When the numbers were measured, so a stale account can say how old
  # they are.
  result["fetchedAtMs"] = 0

  # A panel that is opened and shut repeatedly must not turn into a request
  # per flick, so recent probe results are reused for a short window — and
  # kept as the answer of record when a later probe fails. Every account
  # keeps its own, so one account's allowance never stands in for another's.
  probe_cache = cache_root() / cache_name
  cached = read_json(probe_cache) or {}
  fallback = usable_cached_limits(cached)
  result["fetchedAtMs"] = number(cached.get("fetchedAtMs"))

  # Probing needs a live token and only the Claude Code CLI can mint one: it
  # refreshes the credential file when it runs, so a machine left alone long
  # enough finds the saved token lapsed. Say so — an empty limits list with
  # nothing else set hides the whole section and explains nothing — and keep
  # showing the last numbers, resetting elapsed windows to zero.
  if access_token == "":
    result["limits"] = fallback
    result["usageStatusText"] = "Waiting for auth"
    return result
  if expires_at_ms > 0 and expires_at_ms <= time.time() * 1000:
    result["limits"] = fallback
    # A lapsed access token is routine, not a sign-out: it lives about eight
    # hours, and the CLI mints a new one from the refresh token the next time
    # it runs. Only a lapsed *refresh* token is a sign-in a person must redo,
    # and saying so either way sends people to `claude auth login` nightly --
    # which does not restore anything the CLI would not have restored itself.
    if refresh_expires_at_ms > time.time() * 1000:
      result["usageStatusText"] = "Limits paused"
      result["authHelpText"] = (
        "Claude Code's access token has lapsed"
        + (" — showing the last known limits." if fallback else ".")
        + " Start Claude Code to refresh it; signing in again is not needed."
      )
      return result
    result["usageStatusText"] = "Sign-in expired"
    result["authHelpText"] = (
      "Claude Code's saved sign-in expired"
      + (" — showing the last known limits." if fallback else ".")
      + " Start Claude Code, or run `claude auth login`, to refresh it."
    )
    return result

  # --force is a person asking for fresh numbers, so it skips the reuse window
  # entirely; the interval is there to absorb repeated panel opens, not to
  # overrule someone who pressed refresh.
  fetched_at = number(cached.get("fetchedAtMs")) / 1000
  fetched_age = time.time() - fetched_at
  if fallback and not force and 0 <= fetched_age < PROBE_MIN_INTERVAL_SECONDS:
    result["limits"] = fallback
    result["live"] = True
    result["authHelpText"] = ""
    return result

  probe = probe_limits(access_token)
  if probe["ok"]:
    result["limits"] = probe["limits"]
    result["live"] = True
    result["authHelpText"] = ""
    result["fetchedAtMs"] = round(time.time() * 1000)
    write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"]})
    return result

  # The first probe after login often fires before DHCP has handed out a
  # route. Ask the shell to try again sooner than its regular interval.
  if probe.get("transport"):
    result["retryAdvised"] = True
  if fallback:
    result["limits"] = fallback
    fetched_age = time.time() - result["fetchedAtMs"] / 1000
    result["live"] = 0 <= fetched_age < CURRENT_ENOUGH_SECONDS
  else:
    result["usageStatusText"] = "Claude limits unavailable"
    result["authHelpText"] = probe["helpText"]
  return result


# ---------------------------------------------------------------- accounts


# The accounts `omarchy agent account` registered, read straight from its
# registry. Only a registry holding a second account matters: with one, the
# record is exactly what it always was.
def registered_accounts() -> list[dict[str, Any]]:
  state = Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local" / "state")
  try:
    registry = json.loads((state / "omarchy" / "agents" / "accounts" / "claude.json").read_text(encoding="utf-8"))
  except Exception:
    return []
  accounts = [a for a in registry.get("accounts") or [] if isinstance(a, dict) and a.get("id")]
  if len(accounts) < 2:
    return []
  active_id = registry.get("active") or accounts[0]["id"]
  if not any(a["id"] == active_id for a in accounts):
    active_id = accounts[0]["id"]
  for account in accounts:
    account["active"] = account["id"] == active_id
    account["switch"] = {
      "mode": "auto" if registry.get("switch") == "auto" else "manual",
      "threshold": registry.get("threshold") or 95,
    }
  return accounts


def current_account_id(home: Path) -> str:
  try:
    account = json.loads(profile_path(home).read_text(encoding="utf-8")).get("oauthAccount") or {}
  except Exception:
    return ""
  return str(account.get("accountUuid") or "")


def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]:
  home = expand_path(account["home"]) if account.get("home") else expand_path("~/.claude")
  access_token, expires_at_ms, refresh_expires_at_ms, plan = oauth_login(home)
  # Keyed by the subscription, not the label: forget "Work" and add a different
  # "Work", and the new one must never fall back on the old one's allowance.
  # Who the home is signed in as now, not when it was added: a `claude auth
  # login` run in it by hand makes it a different subscription.
  identity = current_account_id(home)
  key = re.sub(r"[^A-Za-z0-9_-]", "", identity or str(account.get("accountId") or "")) or account["id"]
  # The primary home is keyed the same way once it says who it's signed in
  # as, so signing it in to another subscription never inherits the last
  # one's numbers; until then it keeps the single-account cache.
  if account.get("primary") and not identity:
    cache_name = "claude-limits.json"
  else:
    cache_name = f"claude-limits-{key}.json"
  limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, force, cache_name)
  return {
    "id": account["id"],
    "label": str(account.get("label") or account["id"]),
    "email": str(account.get("email") or ""),
    "plan": plan or str(account.get("plan") or ""),
    "active": account["active"],
    "primary": bool(account.get("primary")),
    "limits": limits["limits"],
    "stale": not limits["live"],
    "fetchedAt": limits["fetchedAtMs"],
    "usageStatusText": limits["usageStatusText"],
    "authHelpText": limits["authHelpText"],
    "retryAdvised": bool(limits.get("retryAdvised")),
  }


# -------------------------------------------------------------------- record


def main() -> int:
  parser = argparse.ArgumentParser()
  parser.add_argument("--force", action="store_true", help="rescan transcripts from scratch and re-probe limits, ignoring caches")
  parser.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh")
  parser.add_argument("--cache-seconds", type=float, default=20)
  args = parser.parse_args()

  claude_dir = config_dir()
  scan_age = 0 if args.force else (900 if args.limits_only else args.cache_seconds)
  stats = cached_scan(claude_dir / "projects", scan_age, use_index=not args.force)

  if number(stats.get("totalPrompts")) <= 0:
    fallback = stats_cache_fallback(claude_dir)
    if fallback is not None:
      stats = fallback
    else:
      # No transcripts and no aggregate cache, but history.jsonl alone can
      # still put numbers on today.
      today_prompts, today_sessions = today_prompts_from_history(claude_dir)
      if today_prompts or today_sessions:
        stats = dict(stats, todayPrompts=today_prompts, todaySessions=today_sessions)

  pi_usage = scan_pi_usage(scan_age)
  if pi_usage is not None:
    stats = merge_stats(stats, pi_usage)

  opencode = scan_opencode_usage(scan_age)
  if opencode is not None:
    stats = merge_stats(stats, opencode)

  registered = registered_accounts()
  accounts = [account_limits(account, args.force) for account in registered]
  current = next((a for a in accounts if a["active"]), None)
  if current:
    # The top-level fields keep describing the account new sessions use, so
    # the bar icon and anything else reading one set of limits stays right.
    plan = current["plan"]
    limits = {key: current[key] for key in ("limits", "usageStatusText", "authHelpText", "retryAdvised")}
    stale, fetched_at = current["stale"], current["fetchedAt"]
  else:
    access_token, expires_at_ms, refresh_expires_at_ms, plan = oauth_login(claude_dir)
    limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, args.force)
    stale, fetched_at = not limits["live"], limits["fetchedAtMs"]

  record = {
    "schemaVersion": 1,
    "id": AGENT_ID,
    "name": AGENT_NAME,
    "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(),
    "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0,
    "hasLocalStats": True,
    "tierLabel": plan,
    "usageStatusText": limits["usageStatusText"],
    "authHelpText": limits["authHelpText"],
    "limits": limits["limits"],
    # Numbers kept past a failed probe, and when they were measured, so the
    # panel can tell them apart from ones checked just now.
    "limitsStale": stale,
    "limitsFetchedAt": fetched_at,
  }
  if limits.get("retryAdvised"):
    record["retryAdvised"] = True
  if accounts:
    record["accountSwitch"] = registered[0]["switch"]
    record["accounts"] = accounts
  record.update(stats)
  print(json.dumps(record, separators=(",", ":"), sort_keys=True))
  return 0


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