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

The plan comes from the settings Grok caches from xAI, and the limit from the
credits endpoint the Grok CLI itself reads for its /usage view, asked with
each account's own sign-in. Sessions come from the summary Grok keeps beside
each session, and tokens and prompts from its usage.json, the ledger `grok
usage` prints: one entry per finished turn, with when it ended and its tokens
by model. The agents panel only ever reads the JSON this prints.
"""

import argparse
import json
import os
import re
import tempfile
import time
import urllib.error
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path

AGENT_ID = "grok"
AGENT_NAME = "Grok"
AUTH_HELP = "Start Grok, or run `grok login`, to sign in."
CREDITS_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits"
# Grok's own CLI watches its subscription once a minute; a panel opened and
# shut repeatedly reuses an answer for about that long.
PROBE_REUSE_SECONDS = 60
# Grok's docs give a saved sign-in 30 days; a refresh token older than that
# can't renew the access token, so it no longer counts as signed in.
REFRESH_LIFETIME_SECONDS = 30 * 24 * 3600
# A limits-only refresh needs fresh limits, not a fresh history scan.
LIMITS_ONLY_SCAN_REUSE_SECONDS = 900


def read_json(path):
  try:
    return json.loads(Path(path).read_text(encoding="utf-8"))
  except Exception:
    return None


# A temporary file of its own per write, so overlapping runs can't rename
# each other's away.
def write_json(path, payload):
  path = Path(path)
  path.parent.mkdir(parents=True, exist_ok=True)
  fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent))
  try:
    with os.fdopen(fd, "w", encoding="utf-8") as handle:
      json.dump(payload, handle)
    os.replace(tmp, path)
  except Exception:
    Path(tmp).unlink(missing_ok=True)
    raise


def cache_root():
  return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage"


def parse_time(value):
  try:
    parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00"))
  except ValueError:
    return None
  return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc)


# The home Grok uses when nothing selects an account: GROK_HOME when set,
# as the CLI itself reads it, else ~/.grok.
def default_home():
  return Path(os.environ.get("GROK_HOME") or (Path.home() / ".grok"))


# Every session under <home>/sessions/<folder>/<session>/: its summary says
# when it was last active and how many prompts it had, and its usage.json
# lists each finished turn with its tokens by model. Added accounts link
# their sessions to the primary's, so this covers all of them. Grok counts
# cached input inside inputTokens; the record keeps the two apart, as the
# other collectors do.
def local_stats(home):
  now = datetime.now().astimezone()
  today = now.date().isoformat()
  week = [(now.date() - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)]
  prompts = sessions = today_sessions = today_prompts = today_tokens = 0
  days = set()
  by_day = {day: 0 for day in week}
  today_by_model = {}
  models = {}
  for session in (Path(home) / "sessions").glob("*/*/"):
    summary = read_json(session / "summary.json")
    if isinstance(summary, dict):
      active = parse_time(summary.get("last_active_at"))
      if active:
        day = active.astimezone().date().isoformat()
        sessions += 1
        prompts += int(summary.get("num_messages") or 0)
        days.add(day)
        if day == today:
          today_sessions += 1

    usage = read_json(session / "usage.json")
    turns = usage.get("turns") if isinstance(usage, dict) else None
    for turn in turns if isinstance(turns, list) else []:
      ended = parse_time(turn.get("endedAt")) if isinstance(turn, dict) else None
      if not ended:
        continue
      day = ended.astimezone().date().isoformat()
      days.add(day)
      if day == today:
        today_prompts += 1
      for model, bucket in (turn.get("modelUsage") or {}).items():
        if not isinstance(bucket, dict):
          continue
        # Grok's totalTokens leaves cache writes out; counting them here keeps
        # the day's total equal to the sum of its model's buckets.
        cached = int(bucket.get("cachedReadTokens") or 0)
        uncached = max(0, int(bucket.get("inputTokens") or 0) - cached)
        output = int(bucket.get("outputTokens") or 0)
        created = int(bucket.get("cacheCreationTokens") or 0)
        tokens = uncached + cached + output + created
        totals = models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0})
        totals["inputTokens"] += uncached
        totals["outputTokens"] += output
        totals["cacheReadInputTokens"] += cached
        totals["cacheCreationInputTokens"] += created
        if day in by_day:
          by_day[day] += tokens
        if day == today:
          today_tokens += tokens
          today_by_model[model] = today_by_model.get(model, 0) + tokens

  return {
    "hasLocalStats": sessions > 0,
    "todayPrompts": today_prompts,
    "todaySessions": today_sessions,
    "todayTotalTokens": today_tokens,
    "todayTokensByModel": today_by_model,
    "recentDays": [{"date": day, "messageCount": by_day[day]} for day in week],
    "totalPrompts": prompts,
    "totalSessions": sessions,
    "activeDays": len(days),
    "activeDates": sorted(days),
    "modelUsage": models,
  }


# A scan from before midnight would count yesterday's sessions as today's.
def cached_local_stats(home, max_age):
  cache = cache_root() / "grok-stats.json"
  cached = read_json(cache) or {}
  today = datetime.now().astimezone().date().isoformat()
  if cached.get("home") == str(home) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age:
    return cached["stats"]
  stats = local_stats(home)
  write_json(cache, {"home": str(home), "day": today, "at": time.time(), "stats": stats})
  return stats


# One login per issuer in auth.json; the first is the CLI's own.
def login(home):
  auth = read_json(Path(home) / "auth.json")
  if not isinstance(auth, dict):
    return {}
  return next((v for v in auth.values() if isinstance(v, dict)), {})


def plan(home):
  cached = read_json(Path(home) / "settings_cache.json") or {}
  try:
    settings = json.loads(cached.get("payload") or "{}").get("settings") or {}
  except Exception:
    return ""
  return str(settings.get("subscription_tier_display") or "")


# The credits answer as one limit window: how much of this period's included
# usage is gone, and when the period ends. The answer is protobuf JSON, which
# leaves out a field holding zero, so a period with nothing used yet has no
# percentage at all.
def credits_limit(data):
  config = data["config"]
  period = config["currentPeriod"]
  return {
    "label": "Weekly" if "WEEK" in str(period.get("type")) else "Monthly",
    "percent": max(0.0, min(1.0, float(config.get("creditUsagePercent") or 0) / 100)),
    "resetsAt": parse_time(period["end"]).isoformat(),
  }


# A window whose reset has passed starts over empty. A weekly one's next
# reset is a whole number of weeks later; a monthly one's isn't known.
def rolled_forward(limits):
  now = datetime.now(timezone.utc)
  rolled = []
  for limit in limits:
    resets = parse_time(limit.get("resetsAt"))
    if resets and resets <= now:
      limit = dict(limit, percent=0.0, resetsAt="")
      if limit.get("label") == "Weekly":
        weeks = (now - resets) // timedelta(weeks=1) + 1
        limit["resetsAt"] = (resets + weeks * timedelta(weeks=1)).isoformat()
    rolled.append(limit)
  return rolled


def collect_limits(home, key, force):
  result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False, "fetchedAtMs": 0}
  entry = login(home)
  token = str(entry.get("key") or "")
  if not token:
    result["usageStatusText"] = "Waiting for auth"
    return result

  cache = cache_root() / f"grok-limits-{key}.json"
  cached = read_json(cache) or {}
  fetched_at = float(cached.get("fetchedAtMs") or 0)
  result["fetchedAtMs"] = fetched_at

  expires = parse_time(entry.get("expires_at"))
  if expires and expires.timestamp() <= time.time():
    result["limits"] = cached.get("limits") or []
    # The access token lives six hours and Grok mints a new one from its
    # refresh token whenever it starts, so a lapsed one is routine, not a
    # sign-out. And while Grok hasn't run, nothing here has spent its
    # allowance, so the last numbers still stand, rolled past any reset.
    if entry.get("refresh_token") and time.time() - expires.timestamp() < REFRESH_LIFETIME_SECONDS:
      result.update(limits=rolled_forward(result["limits"]), live=True)
      if not result["limits"]:
        result["usageStatusText"] = "Limits paused"
        result["authHelpText"] = "Start Grok to see its limits; signing in again is not needed."
    else:
      result["usageStatusText"] = "Sign-in expired"
      result["authHelpText"] = "Grok's saved sign-in expired. Start Grok to refresh it."
    return result

  if cached.get("limits") and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS:
    result.update(limits=cached["limits"], live=True, authHelpText="")
    return result

  request = urllib.request.Request(CREDITS_URL, headers={"Authorization": f"Bearer {token}", "Accept": "application/json"})
  try:
    with urllib.request.urlopen(request, timeout=8) as response:
      limit = credits_limit(json.loads(response.read()))
  except urllib.error.HTTPError as error:
    limit = None
    if error.code in (401, 403):
      result["usageStatusText"] = "Waiting for auth"
  except Exception:
    limit = None

  if limit:
    result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000), authHelpText="")
    write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": result["limits"]})
  else:
    result["limits"] = cached.get("limits") or []
    if not result["usageStatusText"] and not result["limits"]:
      result["usageStatusText"] = "Grok limits unavailable"
  return result


# The accounts `omarchy agent account` registered, read straight from its
# registry. Only a registry holding a second account matters: with one, the
# record describes ~/.grok alone.
def registered_accounts():
  state = Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state"))
  registry = read_json(state / "omarchy" / "agents" / "accounts" / "grok.json") or {}
  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 home_of(account):
  return Path(account.get("home") or (Path.home() / ".grok"))


# Keyed by who the home is signed in as, so a home signed in to someone else
# never shows the last one's numbers.
def cache_key(home, fallback):
  user = str(login(home).get("user_id") or "")
  return re.sub(r"[^A-Za-z0-9_-]", "", user) or fallback


def account_limits(account, force):
  home = home_of(account)
  limits = collect_limits(home, cache_key(home, account["id"]), force)
  return {
    "id": account["id"],
    "label": str(account.get("label") or account["id"]),
    "email": str(account.get("email") or ""),
    "plan": plan(home) 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"],
  }


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument("--force", action="store_true")
  parser.add_argument("--limits-only", action="store_true")
  args = parser.parse_args()

  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:
    tier = current["plan"]
    limits = current
    stale, fetched_at = current["stale"], current["fetchedAt"]
  else:
    home = default_home()
    tier = plan(home)
    limits = collect_limits(home, cache_key(home, "main"), args.force)
    stale, fetched_at = not limits["live"], limits["fetchedAtMs"]

  # Nobody signed in to Grok here: an empty record, which the panel skips.
  # With accounts registered, the primary is ~/.grok by definition.
  stats_home = Path.home() / ".grok" if accounts else default_home()
  signed_in = bool(accounts) or bool(login(stats_home))
  record = {
    "schemaVersion": 1,
    "id": AGENT_ID,
    "name": AGENT_NAME,
    "updatedAt": datetime.now(timezone.utc).isoformat(),
    "ready": signed_in,
    "tierLabel": tier,
    "usageStatusText": limits["usageStatusText"],
    "authHelpText": limits["authHelpText"],
    "limits": limits["limits"],
    "limitsStale": stale,
    "limitsFetchedAt": fetched_at,
  }
  record.update(cached_local_stats(stats_home, LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0))
  if accounts:
    record["accountSwitch"] = registered[0]["switch"]
    record["accounts"] = accounts
  print(json.dumps(record, separators=(",", ":")))


if __name__ == "__main__":
  main()