#!/usr/bin/python3
# blob:summary=Print the Claude Code usage record as JSON
# blob:args=[--force] [--limits-only]
# blob:hidden=true

from __future__ import annotations

import argparse
import datetime as dt
import fcntl
import hashlib
import json
import os
import re
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 usage limits."
USAGE_ENDPOINT = "https://api.anthropic.com/api/oauth/usage"
PROBE_MIN_INTERVAL_SECONDS = 15
CURRENT_ENOUGH_SECONDS = 900
INDEX_VERSION = 4
INDEX_HEAD_BYTES = 4096


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")) / "blob" / "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()

  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,
  }


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:
    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:
  handle_fd, temporary_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp")
  temporary_path = Path(temporary_name)
  try:
    with os.fdopen(handle_fd, "w", encoding="utf-8") as handle:
      handle.write(json.dumps(payload, separators=(",", ":"), sort_keys=True) + "\n")
    temporary_path.chmod(0o644)
    temporary_path.replace(path)
  except BaseException:
    temporary_path.unlink(missing_ok=True)
    raise


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]:
  data = read_json(index_file)
  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"]
  return {}


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]:
  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

  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

  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"):
        break
      offset += len(raw)
      line_number += 1

      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) -> dict[str, Any]:
  previous = load_index(index_file)
  current: dict[str, Any] = {}
  order: list[str] = []
  unreadable_count = 0

  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:
      unreadable_count += 1
      if key in previous:
        current[key] = previous[key]
    if key in current:
      order.append(key)

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

  if current != previous:
    try:
      write_json(index_file, {"version": INDEX_VERSION, "zone": index_zone(), "files": current})
    except Exception as error:
      print(f"Could not write Claude scan index {index_file}: {error}", 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}

  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:
      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),
    "activeDays": len(active_days),
  }


def cached_scan(projects_path: Path, max_age_seconds: float, use_index: bool) -> dict[str, Any]:
  digest = hashlib.sha1(str(projects_path).encode("utf-8")).hexdigest()[:16]
  root = cache_root()
  cache_file = root / f"claude-scan-{digest}.json"
  lock_file = root / f"claude-scan-{digest}.lock"
  index_file = root / f"claude-index-{digest}.json"

  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


def stats_cache_fallback(claude_dir: Path) -> dict[str, Any] | None:
  data = read_json(claude_dir / "stats-cache.json")
  if data is None:
    return None

  today = local_date_string()
  today_tokens = {}
  for entry in data.get("dailyModelTokens") or []:
    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 = {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(value) for value 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),
  }


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)


def oauth_login(claude_dir: Path) -> tuple[str, int, int, str]:
  data = read_json(claude_dir / ".credentials.json") or {}
  login = data.get("claudeAiOauth")
  if not isinstance(login, dict):
    return "", 0, 0, ""
  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,
  )


def profile_tier(claude_dir: Path) -> str:
  profile = read_json(profile_path(claude_dir)) or {}
  account = profile.get("oauthAccount")
  if not isinstance(account, dict):
    return ""
  for key in ("userRateLimitTier", "organizationRateLimitTier"):
    tier = str(account.get(key) or "")
    if tier:
      return tier
  return ""


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"


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
  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():
    timestamp = int(raw)
    if timestamp < 1e12:
      timestamp *= 1000
    try:
      return dt.datetime.fromtimestamp(timestamp / 1000, dt.timezone.utc).isoformat()
    except Exception:
      return raw
  try:
    return dt.datetime.fromisoformat(raw.replace("Z", "+00:00")).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


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 ""


def scoped_limits(payload: dict[str, Any], percent_scale: bool) -> list[dict[str, Any]]:
  entries = payload.get("limits")
  if not isinstance(entries, list):
    return []
  scoped: 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
    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
    scoped.append({
      "label": title,
      "title": title,
      "percent": percent,
      "resetsAt": normalize_reset_at(entry.get("resets_at")),
    })
  return scoped


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 ""
      ) + "."
    else:
      help_text = f"Anthropic's usage endpoint returned status {error.code}."
    return {"ok": False, "helpText": help_text}
  except Exception:
    return {
      "ok": False,
      "transport": True,
      "helpText": "Couldn't reach Anthropic's usage endpoint. Retrying shortly.",
    }

  weekly = usage_bucket(payload, "seven_day_oauth_apps") or usage_bucket(payload, "seven_day")
  session = usage_bucket(payload, "five_hour")
  raw_values = [session.get("utilization") if session else None, weekly.get("utilization") if weekly else None]
  entries = payload.get("limits")
  if isinstance(entries, list):
    raw_values += [entry.get("percent") for entry in entries if isinstance(entry, dict)]
  percent_scale = any(parse_utilization(value) >= 1 for value in raw_values)

  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."}
  return {"ok": True, "limits": limits}


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 for entry in entries if isinstance(entry, dict) and limit_window_open(entry, now)]


def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: int, force: bool) -> dict[str, Any]:
  result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False, "fetchedAtMs": 0}

  probe_cache = cache_root() / "claude-limits.json"
  cached = read_json(probe_cache) or {}
  fallback = usable_cached_limits(cached)
  result["fetchedAtMs"] = number(cached.get("fetchedAtMs"))
  now_ms = time.time() * 1000

  if access_token == "":
    result["limits"] = fallback
    result["usageStatusText"] = "Waiting for sign-in"
    return result

  if expires_at_ms > 0 and expires_at_ms <= now_ms:
    result["limits"] = fallback
    shown = " Showing the last known limits." if fallback else ""
    if refresh_expires_at_ms > now_ms:
      result["usageStatusText"] = "Limits paused"
      result["authHelpText"] = "Claude Code's access token has lapsed." + shown + " Start Claude Code to refresh it."
    else:
      result["usageStatusText"] = "Sign-in expired"
      result["authHelpText"] = "Claude Code's saved sign-in expired." + shown + " Run `claude auth login` to sign in again."
    return result

  fetched_age = time.time() - number(cached.get("fetchedAtMs")) / 1000
  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

  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"] = "Limits unavailable"
    result["authHelpText"] = probe["helpText"]
  return result


def local_stats(claude_dir: Path, scan_age: float, force: bool) -> dict[str, Any]:
  stats = cached_scan(claude_dir / "projects", scan_age, use_index=not force)
  if number(stats.get("totalPrompts")) > 0:
    return stats

  fallback = stats_cache_fallback(claude_dir)
  if fallback is not None:
    return fallback

  today_prompts, today_sessions = today_prompts_from_history(claude_dir)
  if today_prompts or today_sessions:
    return dict(stats, todayPrompts=today_prompts, todaySessions=today_sessions)
  return stats


def main() -> int:
  parser = argparse.ArgumentParser()
  parser.add_argument("--force", action="store_true")
  parser.add_argument("--limits-only", action="store_true")
  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 = local_stats(claude_dir, scan_age, args.force)

  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)

  record = {
    "schemaVersion": 1,
    "id": AGENT_ID,
    "name": AGENT_NAME,
    "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(),
    "tierLabel": plan,
    "usageStatusText": limits["usageStatusText"],
    "authHelpText": limits["authHelpText"],
    "limits": limits["limits"],
    "limitsStale": not limits["live"],
    "limitsFetchedAt": limits["fetchedAtMs"],
  }
  if limits.get("retryAdvised"):
    record["retryAdvised"] = True
  record.update(stats)
  print(json.dumps(record, separators=(",", ":"), sort_keys=True))
  return 0


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