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

Tier 1 (any API key): key usage comes from GET /api/v1/key, and the balance
from GET /api/v1/credits, asked with the management key when there is one
(OpenRouter documents it as management-key only, and an inference key it
refuses isn't asked again for a day); token stats come from the local
opencode SQLite DB
(providerID == "openrouter", role == "assistant") and from pi/omp session
files whose assistant messages ran on the openrouter provider, read-only.
Each key keeps its own last good meter and balance, shown stale when a
check fails. The key is OPENROUTER_API_KEY, opencode's or pi/omp's saved
OpenRouter login, or the config file's. Without any, the local stats still
show and the record waits for auth.

Tier 2 (management key): when OPENROUTER_MANAGEMENT_KEY (or the management
key in the config file) is set, the last-30-days token history comes from
GET /api/v1/activity instead of the local DB -- account-global numbers that
cover every machine and client, grouped by day and model. The /key response
carries is_management_key, so a key that cannot do this fails visibly as
403 and the collector falls back to tier 1 automatically.

Refs: omacom/omarchy#8270 (pi on OpenRouter invisible in panel),
omacom/omarchy#10330 (dedicated pi collector proposal).
"""

from __future__ import annotations

import argparse
import math
import os
import sys
import urllib.request
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any

# The shared collector plumbing ships with the agents panel.
sys.dont_write_bytecode = True
sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents"))
from collectors import (  # noqa: E402
  HttpError,
  LimitsCache,
  Stats,
  age_below,
  base_record as shared_base_record,
  cache_root,  # noqa: F401 - the tests find the cached answers through it
  cached_scan,
  config_home,
  data_home,
  dsh_home,
  emit,
  empty_bucket,
  empty_stats,
  kimchi_sessions_root,
  limits_cache_path,
  now_ms,
  number,
  open_json,
  open_limits,
  opencode_db,
  pi_session_roots,
  read_json,
  routed_match,
  routed_sources,
  scan_cache_paths as shared_scan_cache_paths,
  scan_dsh_sessions,
  scan_max_age,
  scan_opencode as shared_scan_opencode,
  scan_pi_sessions,
  scan_routed,
  stamp,
  warn,
  write_json,
)

AGENT_ID = "openrouter"
AGENT_NAME = "OpenRouter"
PROVIDER = "openrouter"
AUTH_HELP = "Set OPENROUTER_API_KEY, or sign in to OpenRouter in opencode, pi or omp. For the 30-day history, add a management key as OPENROUTER_MANAGEMENT_KEY (openrouter.ai/settings/management-keys)."
MANAGEMENT_AUTH_HELP = "Replace OPENROUTER_MANAGEMENT_KEY (or managementKey in ~/.config/omarchy/agents/openrouter.json) with a working one from openrouter.ai/settings/management-keys."
API_BASE = "https://openrouter.ai/api/v1"


class OpenRouterError(Exception):
  pass


def model_name(raw: Any) -> str:
  value = str(raw or "openrouter").rstrip("/").split("/")[-1]
  return value or "openrouter"


# ---------------------------------------------------------------- auth

def read_config() -> dict[str, Any]:
  return read_json(config_home() / "omarchy" / "agents" / "openrouter.json")


def read_opencode_key(path: Path) -> str:
  entry = read_json(path).get("openrouter")
  if not isinstance(entry, dict):
    return ""
  return str(entry.get("key") or "").strip()


def pi_agent_dir() -> Path:
  # PI_CODING_AGENT_DIR moves pi's agent dir, as pi itself reads it.
  configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip()
  return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent"


def read_pi_key(path: Path) -> str:
  # pi/omp keep a key they were given with /login in auth.json. Only a literal
  # key is taken: a `$VAR` or `!command` entry is resolved by pi alone.
  key = read_opencode_key(path)
  return "" if key.startswith(("!", "$")) else key


def api_key(config: dict[str, Any]) -> str:
  # Explicit key or an opencode, pi or omp login wins; the config file is
  # the last resort for headless/refresh-timer runs without any.
  return (
    str(os.environ.get("OPENROUTER_API_KEY", "")).strip()
    or read_opencode_key(data_home() / "opencode" / "auth.json")
    or read_pi_key(pi_agent_dir() / "auth.json")
    or read_pi_key(Path.home() / ".omp" / "agent" / "auth.json")
    or str(config.get("apiKey") or "").strip()
  )


def management_key(config: dict[str, Any]) -> str:
  # A management key unlocks /activity (30-day account history). A normal
  # inference key answers /key + /credits only and fails /activity with
  # 403 -- never send it there speculatively, just skip tier 2.
  return (
    str(os.environ.get("OPENROUTER_MANAGEMENT_KEY", "")).strip()
    or str(config.get("managementKey") or "").strip()
  )


# ---------------------------------------------------------------- API

# No server answered at all (DNS, connection, timeout): the one failure worth
# a sooner retry.
class OpenRouterTransportError(OpenRouterError):
  pass


# The key was refused. 403 from /activity with an inference key is the
# expected tier-1 signal ("Only management keys can perform this operation"),
# not a broken setup -- callers decide what a refusal means.
class OpenRouterAuthError(OpenRouterError):
  def __init__(self, code: int):
    super().__init__(f"OpenRouter refused the key (status {code})")
    self.code = code


def api_get(path: str, key: str) -> dict:
  req = urllib.request.Request(
    API_BASE + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"}
  )
  try:
    payload = open_json(req, timeout=15)
  except HttpError as error:
    if error.status in (401, 403):
      raise OpenRouterAuthError(error.status) from error
    if error.status:
      raise OpenRouterError(f"OpenRouter's API returned status {error.status}.") from error
    if error.transport:
      # A server that stalls on its headers or body times out the read
      # itself; that is a network failure like any other.
      raise OpenRouterTransportError("Couldn't reach the OpenRouter API. Retrying shortly.") from error
    raise OpenRouterError("OpenRouter returned an invalid response.") from error
  return payload if isinstance(payload, dict) else {}


# /activity names its day as a date, sometimes with a time after it. Anything
# that isn't one files under no day rather than a made-up one.
def activity_day(value: Any) -> str:
  text = str(value or "").strip()[:10]
  try:
    return datetime.strptime(text, "%Y-%m-%d").strftime("%Y-%m-%d")
  except ValueError:
    return ""


def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]:
  """Fold /activity rows into the panel's stats dict.

  One row per (date, endpoint): prompt tokens count as input, and completion
  tokens, reasoning among them, as output (thinking is generated, so it is
  output, as the collector contract has it). The billing API reports
  tokens, never prompt or session counts, so hasPromptStats stays False
  and the panel keeps those numbers out of today's tooltip.
  """
  today_date = today.strftime("%Y-%m-%d")
  recent_dates = [(today - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)]
  recent = {day: 0 for day in recent_dates}
  today_by_model: dict[str, int] = {}
  model_usage: dict[str, dict[str, int]] = {}
  active_dates: set[str] = set()

  for raw in items:
    if not isinstance(raw, dict):
      continue
    day = activity_day(raw.get("date"))
    model = model_name(str(raw.get("model_permaslug") or raw.get("model") or ""))
    prompt = number(raw.get("prompt_tokens"))
    # Reasoning is part of the completion count, as in OpenRouter's usage
    # (its own AI SDK provider reads text = completion - reasoning), so it
    # is never added again; a row naming reasoning alone still counts it.
    output = max(number(raw.get("completion_tokens")), number(raw.get("reasoning_tokens")))
    total = prompt + output
    if total <= 0:
      continue
    bucket = model_usage.setdefault(model, empty_bucket())
    bucket["inputTokens"] += prompt
    bucket["outputTokens"] += output
    if day:
      active_dates.add(day)
    if day in recent:
      recent[day] += total
    if day == today_date:
      today_by_model[model] = today_by_model.get(model, 0) + total

  return {
    "todayTokensByModel": today_by_model,
    "todayTotalTokens": sum(today_by_model.values()),
    "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates],
    "activeDays": len(active_dates),
    "activeDates": sorted(active_dates),
    "modelUsage": model_usage,
  }


def activity_stats(key: str, today: datetime, max_age: float) -> tuple[dict[str, Any] | None, bool]:
  """Tier 2: 30-day account history via management key, or None.

  Returns (stats, transport_failed). Stats are None when there is no
  management key (plain tier-1 run) or the key is refused -- the caller then
  falls back to the local DB scan. The answer is a stats source like the
  local scan, so it is reused on the same schedule; when /activity can't be
  reached, today's last answer stands in for it.
  """
  if not key:
    return None, False
  scan_date = today.strftime("%Y-%m-%d")
  try:
    cache = limits_cache_path(AGENT_ID, key, kind="activity")
    cached = read_json(cache)
  except OSError:
    cache, cached = None, {}
  kept = cached.get("stats") if cached.get("scanDate") == scan_date and isinstance(cached.get("stats"), dict) else None
  if kept is not None and age_below(stamp(cached.get("fetchedAtMs")) / 1000, max_age):
    return kept, False

  try:
    payload = api_get("/activity", key)
  except OpenRouterAuthError:
    return None, False
  except OpenRouterError as error:
    warn(f"activity lookup failed ({error})")
    return kept, isinstance(error, OpenRouterTransportError)
  items = payload.get("data")
  if not isinstance(items, list):
    return kept, False
  stats = empty_stats()
  stats.update(summarize_activity(items, today))
  if cache is not None:
    try:
      write_json(cache, {"schemaVersion": 1, "scanDate": scan_date, "fetchedAtMs": now_ms(), "stats": stats})
    except OSError as exc:
      warn(f"could not write activity cache ({exc})")
  return stats, False


# ---------------------------------------------------------------- local stats

# pi, omp, and opencode name the provider alike. Exact match: a custom
# gateway named like OpenRouter is not this account.
def billed_here(provider: str, model: str, entry: dict) -> bool:
  return provider == PROVIDER


def scan_opencode(stats: Stats) -> bool:
  """Add opencode messages on the openrouter provider, from both of its
  stores. Returns whether the scan ran to completion: a scan cut short by a
  database error still contributes what it read, but must not be cached as
  if it were the whole story."""
  # The closing quote makes the database's prefilter match the whole id.
  rows, complete = shared_scan_opencode([PROVIDER + '"'], billed_here)
  for row in rows:
    stats.add_row(row, model_name(row.model))
  return complete


def pi_roots() -> list[Path]:
  # Kimchi's harness keeps pi-format sessions of its own, and one run with
  # `--provider openrouter` is OpenRouter's.
  return [*pi_session_roots(), kimchi_sessions_root()]


def scan_pi(stats: Stats) -> bool:
  """Add pi/omp and Kimchi messages on the openrouter provider. Returns
  whether every session file could be read, like scan_opencode."""
  rows, complete = scan_pi_sessions(r'"provider"\s*:\s*"openrouter"', billed_here, pi_roots())
  for row in rows:
    stats.add_row(row, model_name(row.model))
  return complete


def scan_dsh(stats: Stats) -> bool:
  """Add DeepSeek Harness attempts on its pi-ai adapter's openrouter route."""
  rows, complete = scan_dsh_sessions(lambda provider, model, request: provider == PROVIDER)
  for row in rows:
    stats.add_row(row, model_name(row.model))
  return complete


def scan_local(stats: Stats, reuse: bool = True) -> bool:
  # Every source always runs, so one failing still leaves the others' numbers.
  opencode_complete = scan_opencode(stats)
  pi_complete = scan_pi(stats)
  dsh_complete = scan_dsh(stats)
  # Claude Code pointed at OpenRouter's Anthropic-compatible API, and MiniMax
  # Code's custom providers (its OpenRouter preset among them) at OpenRouter.
  routed_rows, routed_complete = scan_routed(AGENT_ID, reuse)
  for row in routed_rows:
    stats.add_row(row, model_name(row.model))
  return opencode_complete and pi_complete and dsh_complete and routed_complete


def scan_sources() -> list[Path]:
  return [opencode_db(), *pi_roots(), dsh_home(), *routed_sources()]


def scan_cache_paths() -> tuple[Path, Path]:
  return shared_scan_cache_paths(AGENT_ID, scan_sources())


def cached_local_stats(max_age: float) -> tuple[dict, bool]:
  """Local stats, and whether they cover every source.

  The cache is a pure optimization, and a cached scan was a complete one.
  Concurrent runs (the update command and the panel's own refresh) wait for
  one scan instead of each reading the whole database.
  """

  def scan() -> tuple[dict[str, Any], bool]:
    stats = Stats()
    # A forced scan reads every Codex rollout again too.
    complete = scan_local(stats, reuse=max_age > 0)
    return stats.snapshot(), complete

  return cached_scan(AGENT_ID, scan_sources(), max_age, scan, match=routed_match(), with_complete=True)


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

def base_record(**overrides: Any) -> dict:
  return shared_base_record(AGENT_ID, AGENT_NAME, **overrides)


# A key with a credit limit drains toward empty. limit_reset names how often
# the limit starts over (daily, weekly, monthly), not when, so it labels the
# meter rather than giving it a reset time.
# How long an inference key /credits refused goes unasked: OpenRouter
# documents the endpoint as management-key only, so it answers the same
# until the key's permissions change.
CREDITS_REFUSAL_SECONDS = 24 * 3600

KEY_LIMIT_PERIODS = {"daily": "Daily", "weekly": "Weekly", "monthly": "Monthly"}


# OpenRouter answers `limit: null` for a key without a credit limit, which has
# no meter. An answer missing the limit, or one naming a limit without a
# readable remainder, says nothing about the allowance, so it is a failed check
# that keeps the last good meter rather than an empty one that would replace it.
def key_limit_meter(key_info: dict) -> list[dict]:
  if "limit" not in key_info:
    raise OpenRouterError("OpenRouter answered /key without a credit limit.")
  limit = key_info.get("limit")
  if limit is None:
    return []
  remaining = key_info.get("limit_remaining")
  try:
    if isinstance(limit, bool) or isinstance(remaining, bool):
      raise TypeError
    lim = float(limit)
    rem = max(0.0, float(remaining))
  except (TypeError, ValueError):
    raise OpenRouterError("OpenRouter answered /key with an unreadable credit limit.") from None
  if not (math.isfinite(lim) and math.isfinite(rem)):
    raise OpenRouterError("OpenRouter answered /key with an unreadable credit limit.")
  if lim <= 0:
    return []
  period = KEY_LIMIT_PERIODS.get(str(key_info.get("limit_reset") or "").strip().lower())
  label = f"{period} key limit" if period else "Key limit"
  return [{"label": label, "title": label, "percent": min(1.0, (lim - min(rem, lim)) / lim), "resetsAt": ""}]


def dollars(value: Any) -> float | None:
  if isinstance(value, bool):
    return None
  try:
    parsed = float(value or 0)
  except (TypeError, ValueError):
    return None
  return parsed if math.isfinite(parsed) else None


# Balance: total purchased minus total used, or None for an account that never
# bought credits. An answer missing either total says nothing about the
# balance, so it is a failed check that keeps the last good one, not an empty
# balance that would replace it.
def credits_balance(credits: dict) -> dict | None:
  total_credits = dollars(credits.get("total_credits")) if credits.get("total_credits") is not None else None
  total_usage = dollars(credits.get("total_usage")) if credits.get("total_usage") is not None else None
  if total_credits is None or total_usage is None:
    raise OpenRouterError("OpenRouter answered /credits without a balance.")
  if total_credits <= 0:
    return None
  return {
    "remaining": round(max(0.0, total_credits - total_usage), 6),
    "funded": round(total_credits, 6),
    "spent": round(max(0.0, total_usage), 6),
    "currency": "USD",
    "estimated": False,
  }


# Each credential keeps its own last good answer, so another key never shows
# this one's numbers: the key's meter as openrouter-limits-<key>, the
# balance beside it as openrouter-credits-<key>.
def probe(endpoint: str, key: str, kind: str, parse, force: bool) -> tuple[dict | None, OpenRouterError | None]:
  """One endpoint's answer as (part, error).

  Opening and shutting the panel must not turn into three API calls per
  flick, so a probe this recent is reused; --force, a person pressing
  refresh, skips it. A fresh answer replaces the cached one, and a failed one
  falls back on it. part is None when nothing was ever fetched.
  A key with no credit limit has an empty meter, and its answer is reused
  all the same.
  """
  cache = LimitsCache(AGENT_ID, key, kind=kind)
  kept = dict(cache.data, fetchedAtMs=cache.fetched_at) if cache.fetched_at else None
  if kept is not None and cache.fresh(force, allow_empty=True):
    return kept, None
  try:
    payload = api_get(endpoint, key)
    data = payload.get("data")
    part = parse(data if isinstance(data, dict) else {})
  except OpenRouterError as error:
    return kept, error
  # Refreshes can overlap (the panel's own and a sooner retry), so an answer
  # only replaces one fetched before it.
  cache.save(part.pop("limits", []), only_if_newer=True, **part)
  return dict(cache.data), None


def collect_billing(key: str, manager: str, force: bool) -> dict[str, Any]:
  """The key-limit meter and the balance, each kept past a failed check.

  The meter describes an inference key's own credit limit, so with only a
  management key there is none to show; the balance is the account's, and
  an inference key reads it only where OpenRouter still lets it.
  """
  result: dict[str, Any] = {
    "limits": [], "limitsStale": False, "limitsFetchedAt": 0,
    "usageStatusText": "", "authHelpText": "", "balance": None, "retryAdvised": False,
  }
  parts: dict[str, dict | None] = {}
  errors: dict[str, OpenRouterError] = {}
  if key:
    parts["key"], error = probe("/key", key, "limits", lambda data: {"limits": key_limit_meter(data)}, force)
    if error:
      errors["key"] = error
  # OpenRouter documents /credits as management-key only, answering another
  # key 403 "Only management keys can perform this operation". So the
  # management key asks when there is one, and an inference key alone asks
  # until it is refused that way, which is no balance for this key, not a
  # refused key: /key says whether the key itself still works. The refusal
  # is kept, so the next refreshes don't ask again. A management key that is
  # refused itself (revoked, or no management key at all) leaves the
  # inference key to ask; when that key has no balance either, the
  # management key's refusal stands, with its last good balance kept.
  error = manager_error = None
  if manager:
    parts["credits"], error = probe("/credits", manager, "credits", lambda data: {"balance": credits_balance(data)}, force)
    manager_error = error
  if key and (not manager or isinstance(error, OpenRouterAuthError)):
    managed = parts.get("credits"), error
    refusal = LimitsCache(AGENT_ID, key, kind="credits")
    if refusal.data.get("refused") is True and refusal.younger_than(CREDITS_REFUSAL_SECONDS):
      parts["credits"], error = managed
    else:
      parts["credits"], error = probe("/credits", key, "credits", lambda data: {"balance": credits_balance(data)}, force)
      if isinstance(error, OpenRouterAuthError) and error.code == 403:
        refusal.save([], refused=True)
        parts["credits"], error = managed
  if error:
    errors["credits"] = error

  meter = parts.get("key")
  if meter is not None:
    # A kept window whose reset has passed is an untouched allowance, not an
    # unknown one. The key meter only ever names a period, so this guards a
    # cache written by a meter that carries a reset time.
    result["limits"] = open_limits(meter.get("limits"))
  credits = parts.get("credits")
  if credits is not None and isinstance(credits.get("balance"), dict):
    result["balance"] = credits["balance"]
  kept = [part for part in parts.values() if part is not None]
  if kept:
    result["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in kept)
  # Anything that didn't come back fresh marks the record stale, kept or not;
  # the panel shows the kept numbers' age on hover.
  result["limitsStale"] = bool(errors)
  result["retryAdvised"] = any(isinstance(error, OpenRouterTransportError) for error in errors.values())

  refused = [error for error in errors.values() if isinstance(error, OpenRouterAuthError)]
  if refused and refused == [manager_error]:
    # Only the management key was refused; the inference key, if any, still
    # works, so the help names the key to replace.
    result["usageStatusText"] = "Sign-in expired"
    result["authHelpText"] = "OpenRouter refused the management key. " + MANAGEMENT_AUTH_HELP
  elif refused:
    result["usageStatusText"] = "Sign-in expired"
    result["authHelpText"] = "OpenRouter refused the API key. " + AUTH_HELP
  elif "key" in errors and meter is None:
    result["usageStatusText"] = "OpenRouter limits unavailable"
    result["authHelpText"] = str(errors["key"])
  elif "credits" in errors and credits is None:
    result["usageStatusText"] = "Balance unavailable"
    result["authHelpText"] = str(errors["credits"])
  elif errors:
    result["authHelpText"] = str(next(iter(errors.values())))
  return result


def scan(args: Any) -> dict[str, Any]:
  config = read_config()
  key = api_key(config)
  manager = management_key(config)
  today = datetime.now().astimezone()
  max_age = scan_max_age(args.force, args.limits_only)

  if not key and not manager:
    # No key anywhere: there is no balance or meter to read, but what pi,
    # omp and opencode spent through OpenRouter here is still OpenRouter's
    # usage, and no other collector counts it. With none, OpenRouter isn't
    # used here and the panel skips it.
    stats, complete = cached_local_stats(max_age)
    if number(stats.get("totalPrompts")) <= 0 and not stats.get("modelUsage"):
      return base_record()
    record = base_record(ready=True, hasLocalStats=complete, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP)
    record.update(stats)
    return record

  # Tier 2 first: with a management key the 30-day account history covers
  # every machine and client. Without one (or a refused one) fall back to
  # the local opencode and pi scan below.
  managed, activity_unreachable = activity_stats(manager, today, max_age)
  if managed is not None:
    stats = dict(managed)
    complete = True
    has_prompt_stats = False
    # Tier-2 numbers are account-global (billing API), like Fireworks: every
    # synced device reports the same truth, so aggregation must merge them by
    # widest value, not sum them. Tier 1 is machine-local instead.
    stats["scope"] = "account"
  else:
    stats, complete = cached_local_stats(max_age)
    has_prompt_stats = True
  has_data = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage"))

  billing = collect_billing(key, manager, args.force)

  record = base_record(
    ready=True,
    tierLabel="Prepaid",
    hasLocalStats=complete,
    hasPromptStats=has_prompt_stats,
    limits=billing["limits"],
    limitsStale=billing["limitsStale"],
    limitsFetchedAt=billing["limitsFetchedAt"],
    usageStatusText=billing["usageStatusText"],
    authHelpText=billing["authHelpText"],
  )
  # Usually the seconds after login before the network is up; the panel
  # tries once more soon instead of waiting out its refresh interval.
  if billing["retryAdvised"] or activity_unreachable:
    record["retryAdvised"] = True
  record.update(stats)
  if billing["balance"]:
    record["balance"] = billing["balance"]
  # A key nothing has been spent through, with no credits bought and no limit
  # set, has nothing to show: the machine doesn't use OpenRouter yet.
  if not has_data and not billing["balance"] and not record["limits"] and not record["usageStatusText"] and not record["limitsStale"]:
    record["ready"] = False
    record["tierLabel"] = ""
  return record


def main() -> int:
  ap = argparse.ArgumentParser(description="Print the OpenRouter usage record as JSON")
  # --force rescans and re-probes everything; --limits-only reuses a recent
  # local scan or activity history for up to 15 minutes, while the key meter
  # and balance are probed like a normal run.
  ap.add_argument("--force", action="store_true", help="rescan sessions and re-probe the API, ignoring caches")
  ap.add_argument("--limits-only", action="store_true", help="reuse any recent session scan; only the limits probe must be fresh")
  args = ap.parse_args()
  try:
    record = scan(args)
  except Exception as error:
    record = base_record(usageStatusText="OpenRouter unavailable", authHelpText="OpenRouter usage scan failed")
    warn(type(error).__name__)
  emit(record)
  return 0


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