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

Local stats come from native Codex CLI session files (and pi sessions that
ran through openai-codex); rate limits and the plan come from the Codex
app-server RPC. The agents panel only ever reads the JSON this prints.
"""

import argparse
import json
import os
import select
import shutil
import subprocess
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

AGENT_ID = "codex"
AGENT_NAME = "Codex"
AUTH_HELP = "Run `codex login` to authenticate."


def local_day(value):
  if value is None:
    return datetime.now().strftime("%Y-%m-%d")
  if isinstance(value, (int, float)):
    # pi message timestamps are milliseconds; Codex timestamps are usually seconds.
    if value > 10_000_000_000:
      value = value / 1000
    return datetime.fromtimestamp(value).strftime("%Y-%m-%d")
  text = str(value)
  try:
    if text.endswith("Z"):
      dt = datetime.fromisoformat(text[:-1] + "+00:00")
    else:
      dt = datetime.fromisoformat(text)
    if dt.tzinfo is not None:
      dt = dt.astimezone()
    return dt.strftime("%Y-%m-%d")
  except Exception:
    return datetime.now().strftime("%Y-%m-%d")


def number(value):
  try:
    return int(value or 0)
  except Exception:
    return 0


def model_name(raw):
  value = str(raw or "codex")
  return value if value else "codex"


def runtime_env():
  home = str(Path.home())
  path_parts = [
    os.environ.get("PATH", ""),
    f"{home}/.local/bin",
    f"{home}/.npm-global/bin",
    f"{home}/.local/share/mise/shims",
  ]
  env = os.environ.copy()
  env["PATH"] = os.pathsep.join(part for part in path_parts if part)
  return env


ENV = runtime_env()


def find_command(name):
  return shutil.which(name, path=ENV.get("PATH"))


now = datetime.now()
today = now.strftime("%Y-%m-%d")
recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)]
recent = {day: {"date": day, "messageCount": 0} for day in recent_dates}
today_tokens_by_model = {}
model_usage = {}
today_sessions = set()
active_days = set()

today_prompts = 0
today_total_tokens = 0
total_prompts = 0
total_sessions = set()
seen_pi_messages = set()


def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write):
  global today_prompts, today_total_tokens, total_prompts
  total = input_tokens + output_tokens + cache_read + cache_write
  total_prompts += 1
  total_sessions.add(session_key)
  active_days.add(day)

  bucket = model_usage.setdefault(model, {
    "inputTokens": 0,
    "outputTokens": 0,
    "cacheReadInputTokens": 0,
    "cacheCreationInputTokens": 0,
  })
  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_prompts += 1
    today_sessions.add(session_key)
    today_total_tokens += total
    today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total


def scan_pi_sessions():
  root = Path.home() / ".pi" / "agent" / "sessions"
  if not root.exists():
    return
  try:
    rg = find_command("rg") or "rg"
    proc = subprocess.Popen(
      [rg, "--json", "-e", '"provider":"openai-codex"', "-e", '"api":"openai-codex"', str(root)],
      stdout=subprocess.PIPE,
      stderr=subprocess.DEVNULL,
      text=True,
      errors="replace",
      env=ENV,
    )
  except FileNotFoundError:
    return

  assert proc.stdout is not None
  for raw in proc.stdout:
    try:
      event = json.loads(raw)
      if event.get("type") != "match":
        continue
      line = event.get("data", {}).get("lines", {}).get("text", "")
      path = event.get("data", {}).get("path", {}).get("text", "pi-session")
      entry = json.loads(line)
    except Exception:
      continue

    if entry.get("type") != "message":
      continue
    message_key = path + ":" + str(entry.get("id") or "")
    if message_key in seen_pi_messages:
      continue
    seen_pi_messages.add(message_key)
    message = entry.get("message") or {}
    if message.get("role") != "assistant":
      continue
    provider = str(message.get("provider") or "")
    api = str(message.get("api") or "")
    if provider != "openai-codex" and not api.startswith("openai-codex"):
      continue

    usage = message.get("usage") or {}
    if not usage:
      continue
    total = number(usage.get("totalTokens"))
    input_tokens = number(usage.get("input"))
    output_tokens = number(usage.get("output"))
    cache_read = number(usage.get("cacheRead"))
    cache_write = number(usage.get("cacheWrite"))
    if total and not (input_tokens or output_tokens or cache_read or cache_write):
      input_tokens = total
    if not (input_tokens or output_tokens or cache_read or cache_write):
      continue

    day = local_day(entry.get("timestamp") or message.get("timestamp"))
    session_key = path
    add_usage(day, session_key, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write)

  try:
    proc.wait(timeout=1)
  except Exception:
    proc.kill()


def scan_native_codex_sessions():
  codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
  roots = [codex_home / "sessions", codex_home / "archived_sessions"]
  files = []
  cutoff = time.time() - 30 * 24 * 60 * 60
  for root in roots:
    if not root.exists():
      continue
    for path in root.rglob("*.jsonl"):
      try:
        if path.stat().st_mtime >= cutoff:
          files.append(path)
      except OSError:
        pass

  for path in files:
    current_model = "codex"
    try:
      with path.open(errors="replace") as handle:
        for raw in handle:
          try:
            entry = json.loads(raw)
          except Exception:
            continue
          if entry.get("type") == "turn_context":
            payload = entry.get("payload") or {}
            current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model)
            continue
          payload = entry.get("payload") or entry
          if entry.get("type") == "response_item" and isinstance(payload, dict):
            payload = payload.get("payload") or payload
          if not isinstance(payload, dict):
            continue
          if payload.get("type") != "token_count":
            continue
          info = payload.get("info") or {}
          # total_token_usage is cumulative for the session. Adding every
          # snapshot makes usage grow quadratically, so count the last turn.
          usage = info.get("last_token_usage") or {}
          cache_read = number(usage.get("cached_input_tokens"))
          cache_write = number(usage.get("cache_write_input_tokens"))
          # Cached tokens are included in input_tokens, and reasoning tokens
          # are included in output_tokens. Keep the cache split without
          # counting either category twice.
          input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write)
          output_tokens = number(usage.get("output_tokens"))
          if not (input_tokens or output_tokens or cache_read or cache_write):
            continue
          day = local_day(entry.get("timestamp") or path.stat().st_mtime)
          add_usage(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write)
    except Exception:
      continue


def rpc_request(proc, request_id, method, params=None, timeout=8):
  payload = {"id": request_id, "method": method, "params": params or {}}
  proc.stdin.write(json.dumps(payload) + "\n")
  proc.stdin.flush()
  deadline = time.time() + timeout
  while time.time() < deadline:
    ready, _, _ = select.select([proc.stdout], [], [], 0.25)
    if not ready:
      continue
    line = proc.stdout.readline()
    if not line:
      break
    try:
      message = json.loads(line)
    except Exception:
      continue
    if message.get("id") == request_id:
      return message
  raise TimeoutError(method)


def limit_window(window):
  if not isinstance(window, dict):
    return None
  used = window.get("usedPercent")
  if used is None:
    return None
  mins = number(window.get("windowDurationMins"))
  if mins == 10080:
    label = "Weekly (7-day)"
  elif mins and mins % 60 == 0:
    label = f"{mins // 60}h window"
  elif mins:
    label = f"{mins}m window"
  else:
    label = "Limit"
  reset = window.get("resetsAt")
  return {
    "label": label,
    "percent": float(used) / 100.0,
    "resetsAt": datetime.fromtimestamp(number(reset), timezone.utc).isoformat() if reset else "",
  }


def fetch_codex_rpc():
  result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP}
  codex = find_command("codex")
  if not codex:
    result["usageStatusText"] = "Codex unavailable"
    result["authHelpText"] = "codex not found in PATH"
    return result

  try:
    proc = subprocess.Popen(
      [codex, "-s", "read-only", "-a", "untrusted", "app-server"],
      stdin=subprocess.PIPE,
      stdout=subprocess.PIPE,
      stderr=subprocess.DEVNULL,
      text=True,
      env=ENV,
    )
  except Exception as exc:
    result["usageStatusText"] = "Codex unavailable"
    result["authHelpText"] = str(exc)
    return result

  try:
    rpc_request(proc, 1, "initialize", {"clientInfo": {"name": "omarchy-agent-usage", "version": "1"}}, timeout=8)
    proc.stdin.write(json.dumps({"method": "initialized", "params": {}}) + "\n")
    proc.stdin.flush()
    account_msg = rpc_request(proc, 2, "account/read", timeout=4)
    limits_msg = rpc_request(proc, 3, "account/rateLimits/read", timeout=4)

    account = (account_msg.get("result") or {}).get("account") or {}
    limits = (limits_msg.get("result") or {}).get("rateLimits") or {}
    plan = limits.get("planType") or account.get("planType") or account.get("type") or ""
    result["tierLabel"] = str(plan) if plan else ""

    for window in (limits.get("primary"), limits.get("secondary")):
      entry = limit_window(window)
      if entry:
        result["limits"].append(entry)
  except Exception as exc:
    result["usageStatusText"] = "Codex limits unavailable"
    result["authHelpText"] = str(exc)
  finally:
    try:
      proc.terminate()
      proc.wait(timeout=1)
    except Exception:
      try:
        proc.kill()
      except Exception:
        pass
  return result


def main():
  parser = argparse.ArgumentParser()
  # Local stats and limits come from the same cheap scan, so there is no
  # cache to force past and no faster limits-only path. The flags exist so
  # every collector accepts the same invocation.
  parser.add_argument("--force", action="store_true")
  parser.add_argument("--limits-only", action="store_true")
  parser.parse_args()

  scan_pi_sessions()
  scan_native_codex_sessions()
  rpc = fetch_codex_rpc()

  record = {
    "schemaVersion": 1,
    "id": AGENT_ID,
    "name": AGENT_NAME,
    "updatedAt": datetime.now(timezone.utc).isoformat(),
    "ready": True,
    "hasLocalStats": True,
    "todayPrompts": today_prompts,
    "todaySessions": len(today_sessions),
    "todayTotalTokens": today_total_tokens,
    "todayTokensByModel": today_tokens_by_model,
    "recentDays": [recent[day] for day in recent_dates],
    "totalPrompts": total_prompts,
    "totalSessions": len(total_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),
    "modelUsage": model_usage,
  }
  record.update(rpc)
  print(json.dumps(record, separators=(",", ":")))


if __name__ == "__main__":
  main()