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

Local stats come from pi/omp sessions that ran through an opencode provider
and from opencode's own message database, both of its stores. Only the Zen
and Go providers count: opencode runs on Anthropic or OpenAI belong to the
claude and codex records. The limits are the Go plan's rolling, weekly, and
monthly meters from OpenCode's usage endpoint. The agents panel only ever
reads the JSON this prints.
"""

import argparse
import json
import math
import os
import sqlite3
import sys
import time
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import NamedTuple

# 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,
  base_record,
  cache_root,  # noqa: F401 - the tests reach the cache through it
  cached_scan,
  data_home,
  emit,
  limits_cache_path as shared_limits_cache_path,
  number,
  open_json,
  opencode_db,
  pi_session_roots,
  scan_max_age,
  scan_opencode,
  scan_pi_sessions,
)

AGENT_ID = "opencode"
AGENT_NAME = "OpenCode"

# The subscription meters live on the Zen "go" gateway; the plain gateway has
# no usage endpoint. Limits are per-account, so they come from the same key
# regardless of which provider id a given run used. A Console session asks
# Console's own endpoint for the same numbers.
API_BASE_URL = "https://opencode.ai"
USAGE_PATH = "/zen/go/v1/usage"
CONSOLE_USAGE_PATH = "/inference/go/v1/usage"
CONSOLE_ORG_HEADER = "x-opencode-org-id"
CONSOLE_EXPIRED_HELP = "OpenCode's Console session has expired. Run `opencode auth login` to refresh it."

# A key with no Go plan behind it (a Zen pay-as-you-go account) gets the same
# answer every time, so it is asked again hourly rather than on every run;
# --force asks at once, for someone who just subscribed.
NO_PLAN_RECHECK_SECONDS = 3600

# The subscription has two provider ids: "opencode" for the plain Zen
# gateway and "opencode-go" for the Go plan. Both spend the same account and
# both appear in pi sessions and in opencode's own database.
PROVIDERS = ("opencode", "opencode-go")


# Usage percentages arrive as 0-100; the record contract wants 0-1. An
# unreadable value returns -1 so the caller can drop the window.
def fraction(value):
  if isinstance(value, bool):
    return -1.0
  try:
    percent = float(value) / 100.0
  except (TypeError, ValueError):
    return -1.0
  return percent if math.isfinite(percent) else -1.0


def model_name(raw):
  # opencode stores the model as a {providerID, modelID} object in its own
  # database and as a plain string in pi sessions; the shared readers hand
  # back the id either way. A variant suffix such as ":high" names the
  # thinking level, not the model.
  value = str(raw or "").rstrip("/").split("/")[-1].split(":")[0] or "opencode"
  # The panel title-cases words but cannot produce all-caps acronyms.
  if value[:3].lower() == "glm":
    value = "GLM" + value[3:]
  return value


# Exact match: opencode provider ids are free-form, and a custom
# "opencode-local" gateway is not this subscription. Anthropic and OpenAI runs
# never match, so the claude and codex collectors keep those.
def billed_here(provider, model, entry):
  return provider in PROVIDERS


def run_local_scans():
  # pi and omp can spend the subscription without opencode's CLI ever
  # running, so its own database cannot be the only source. Their JSONL
  # transcripts carry the provider, model, and token usage per message. A
  # subscription burned entirely through opencode leaves no pi session files,
  # but opencode records the same per message in its own database.
  stats = Stats()
  pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"opencode(-go)?"', billed_here)
  db_rows, db_complete = scan_opencode(PROVIDERS, billed_here)
  for row in pi_rows + db_rows:
    stats.add_row(row, model_name(row.model))
  return stats.snapshot(), pi_complete and db_complete


def cached_local_stats(max_age):
  # The digest covers every data path the scan reads: the opencode DB and the
  # pi/omp session roots.
  return cached_scan(AGENT_ID, [opencode_db(), *pi_session_roots()], max_age, run_local_scans)


class Credential(NamedTuple):
  """One way of asking OpenCode for the Go meters."""

  token: str
  url: str
  headers: dict
  expired: bool = False


def read_auth_key(path):
  try:
    parsed = json.loads(path.read_text(encoding="utf-8"))
  except Exception:
    # Unreadable, not JSON, or nested past the parser's depth (a
    # RecursionError): no key here.
    return ""
  if not isinstance(parsed, dict):
    return ""
  # Only the subscription's own provider ids: a custom gateway named like
  # "opencode-local" holds a key for somewhere else, never to be sent here.
  # An entry is an {type, key} object, or the key as a plain string.
  # A pi `$VAR` or `!command` entry is resolved by pi alone, never sent.
  entries = [parsed[provider] for provider in PROVIDERS if provider in parsed]
  for entry in entries:
    if isinstance(entry, dict):
      entry = entry.get("key")
    if isinstance(entry, str) and entry.strip() and not entry.strip().startswith(("!", "$")):
      return entry.strip()
  return ""


def pi_agent_dir():
  configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip()
  return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent"


def read_store_credentials(db):
  """The active credentials opencode V2 keeps in its database: Go keys, and
  the Console session it signs the opencode integration in with."""
  keys, console = [], None
  if not db.is_file():
    return keys, console
  try:
    conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2)
  except sqlite3.Error:
    return keys, console
  try:
    conn.execute("PRAGMA query_only = ON")
    rows = conn.execute(
      "SELECT integration_id, value FROM credential WHERE active = 1 AND integration_id IN ('opencode-go', 'opencode')"
      " ORDER BY integration_id = 'opencode-go' DESC, time_updated DESC"
    ).fetchall()
  except sqlite3.Error:
    return keys, console
  finally:
    conn.close()
  for _, raw in rows:
    try:
      value = json.loads(raw)
    except Exception:
      # Not JSON, or nested past the parser's depth: skip this row alone.
      continue
    if not isinstance(value, dict):
      continue
    if value.get("type") == "key" and str(value.get("key") or "").strip():
      keys.append(str(value["key"]).strip())
    elif value.get("type") == "oauth" and console is None and str(value.get("access") or "").strip():
      # Console traffic is scoped to a workspace, and a wrong one is a hard
      # 403, so the header only goes out when the id is there.
      metadata = value.get("metadata")
      org = str(metadata.get("orgID") or "").strip() if isinstance(metadata, dict) else ""
      expired = 0 < number(value.get("expires")) <= time.time() * 1000
      console = Credential(str(value["access"]).strip(), API_BASE_URL + CONSOLE_USAGE_PATH, {CONSOLE_ORG_HEADER: org} if org else {}, expired)
  return keys, console


def credential_candidates():
  # An explicit key wins, then opencode's own credential stores, then pi's:
  # pi drives the same Zen subscription and records it under the same
  # provider ids, so a machine that codes through pi never signs the opencode
  # CLI in. The Console session is the implicit one, so it comes last. Which
  # one the endpoint accepts isn't knowable from here, so a refused one only
  # hands over to the next.
  store_keys, console = read_store_credentials(opencode_db())
  tokens = [str(os.environ.get(name, "")).strip() for name in ("OPENCODE_API_KEY", "OPENCODE_GO_API_KEY")]
  tokens += store_keys
  tokens.append(read_auth_key(data_home() / "opencode" / "auth.json"))
  tokens.append(read_auth_key(pi_agent_dir() / "auth.json"))
  tokens.append(read_auth_key(Path.home() / ".omp" / "agent" / "auth.json"))
  candidates = []
  for token in tokens:
    if token and token not in [c.token for c in candidates]:
      candidates.append(Credential(token, API_BASE_URL + USAGE_PATH, {}))
  if console is not None:
    candidates.append(console)
  return candidates


def reset_from_seconds(value):
  seconds = number(value)
  if seconds <= 0:
    return ""
  try:
    return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat()
  except OverflowError:
    # Further out than a date goes: no reset time to show.
    return ""


def iso_reset(value):
  # Only an ISO time is a reset time; an epoch, a list or garbage is none.
  if not isinstance(value, str) or not value.strip():
    return ""
  try:
    datetime.fromisoformat(value.strip().replace("Z", "+00:00"))
  except ValueError:
    return ""
  return value.strip()


# The live endpoint nests the windows under "usage" with "percent" and
# "resetsAt"; another deployment spells them "<window>Usage" with
# "usagePercent" and "resetInSec". Both are read, so the collector doesn't
# depend on which one answers.
def parse_usage(payload):
  usage = payload.get("usage") if isinstance(payload, dict) else None
  usage = usage if isinstance(usage, dict) else (payload if isinstance(payload, dict) else {})
  limits = []
  for name, label in (("rolling", "Rolling (5h)"), ("weekly", "Weekly (7-day)"), ("monthly", "Monthly")):
    window = usage.get(name)
    if not isinstance(window, dict):
      window = payload.get(name + "Usage") if isinstance(payload, dict) else None
    if not isinstance(window, dict):
      continue
    percent = fraction(window.get("percent", window.get("usagePercent")))
    if percent < 0:
      continue
    # A rate-limited window is spent until it resets, whatever it reads.
    if str(window.get("status") or "") == "rate-limited":
      percent = 1.0
    limits.append({
      "label": label,
      "percent": max(0.0, min(1.0, percent)),
      "resetsAt": iso_reset(window.get("resetsAt")) or reset_from_seconds(window.get("resetInSec")),
    })
  return limits


def probe_limits(credential):
  request = urllib.request.Request(credential.url, headers={
    "Authorization": "Bearer " + credential.token,
    "Accept": "application/json",
    # Cloudflare answers the default urllib signature with a 1010 before the
    # request reaches the API, so identify the collector instead.
    "User-Agent": "omarchy-agent-usage-opencode/1.0",
    **credential.headers,
  })
  unreachable = {"ok": False, "transport": True, "helpText": "Couldn't reach OpenCode's usage endpoint. Retrying shortly. Local OpenCode stats are still shown."}
  try:
    payload = open_json(request, timeout=10)
  except HttpError as error:
    if error.transport:
      return unreachable
    # A refused credential may just be the wrong one for this account, so the
    # next way in still gets its turn.
    if error.status == 401:
      return {"ok": False, "refused": True, "helpText": "OpenCode rejected the saved key. Reconnect OpenCode Go with /connect in the TUI, or set OPENCODE_API_KEY."}
    if error.status == 403:
      # A key without a Go plan behind it: Zen pay-as-you-go has no meters,
      # which is no trouble at all.
      return {"ok": False, "noPlan": True}
    if error.status:
      return {"ok": False, "helpText": f"OpenCode's usage endpoint returned status {error.status}. Local OpenCode stats are still shown."}
    payload = None
  except Exception:
    # open_json already names every transport failure, so anything else is
    # a broken answer, and retrying soon won't mend it.
    payload = None

  if not isinstance(payload, dict) or not payload:
    return {"ok": False, "helpText": "OpenCode's usage endpoint returned no data. Local OpenCode stats are still shown."}
  limits = parse_usage(payload)
  if not limits:
    return {"ok": False, "helpText": "OpenCode's usage endpoint returned no quota windows. Local OpenCode stats are still shown."}
  return {"ok": True, "limits": limits}


# One cache per credential, so a machine signed in to another account never
# shows the last one's meters.
def limits_cache_path(token):
  return shared_limits_cache_path(AGENT_ID, token)


def collect_limits(force):
  result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "stale": False, "fetchedAtMs": 0, "signedIn": False}
  # No key at all is a Zen user paying as they go, or someone who only runs
  # opencode locally: local stats stand alone, with nothing to fix.
  candidates = credential_candidates()
  if not candidates:
    return result
  result["signedIn"] = True

  # Limits only ever show for a Go plan, so a tier follows them.
  def answer(limits, fetched_at, stale):
    # Stale says limits were kept past a failed check, so none kept is none stale.
    result.update(limits=limits, tierLabel="Go" if limits else "", stale=stale and bool(limits), fetchedAtMs=fetched_at if limits else 0)
    return result

  refused = None
  no_plan = False
  for candidate in candidates:
    cache = LimitsCache(AGENT_ID, candidate.token)
    fallback, fetched_at = cache.limits, cache.fetched_at

    # A key already known to have no Go plan isn't asked again for a while.
    if cache.data.get("noPlan") is True and not force and cache.younger_than(NO_PLAN_RECHECK_SECONDS):
      no_plan = True
      continue

    # A recent answer stands in for a probe; --force skips the wait.
    if cache.fresh(force):
      return answer(fallback, fetched_at, False)

    probe = probe_limits(candidate)
    if probe.get("ok"):
      return answer(probe["limits"], cache.save(probe["limits"], only_if_newer=True), False)

    if probe.get("noPlan"):
      # The windows of a plan that's gone mean nothing now, so they go too.
      cache.save([], only_if_newer=True, noPlan=True)
      no_plan = True
      continue

    if probe.get("refused"):
      # Keep the first refusal that still has open windows, so a later
      # credential without a cache can't blank them. A lapsed Console
      # session is what to fix, whichever refusal is kept.
      expired = candidate.expired or bool(refused and refused[3])
      if refused is None or (not refused[1] and fallback):
        refused = (probe, fallback, fetched_at, expired)
      else:
        refused = refused[:3] + (expired,)
      continue

    # Last good numbers, marked stale, beat an empty limits section while
    # the probe is down. A key that already said it has no plan makes an
    # outage on another one no news.
    result["retryAdvised"] = bool(probe.get("transport"))
    answer(fallback, fetched_at, True)
    if fallback:
      result["authHelpText"] = probe["helpText"]
    elif not no_plan:
      result["usageStatusText"] = "OpenCode limits unavailable"
      result["authHelpText"] = probe["helpText"]
    return result

  # A working key without a Go plan is the answer: the other ways in only
  # failed to say the same.
  if no_plan:
    return result

  # Every way in was refused: an account problem, not an outage. Keep any
  # still-open numbers, but say a new sign-in is needed.
  probe, fallback, fetched_at, expired = refused
  answer(fallback, fetched_at, True)
  result["usageStatusText"] = "Sign-in expired"
  result["authHelpText"] = CONSOLE_EXPIRED_HELP if expired else probe["helpText"]
  return result


def main():
  parser = argparse.ArgumentParser()
  # --force rescans everything and rewrites the cache. --limits-only is kept
  # for CLI compatibility with the panel's refreshLimits() call: only the
  # limits probe must be fresh, so it may reuse a scan for far longer than a
  # normal run, whose short window exists purely to dedup concurrent
  # collector runs.
  parser.add_argument("--force", action="store_true")
  parser.add_argument("--limits-only", action="store_true")
  args = parser.parse_args()

  stats = cached_local_stats(scan_max_age(args.force, args.limits_only))
  limits = collect_limits(args.force)

  # Nobody uses OpenCode here: no usage and no credential leave a record the
  # panel skips.
  record = base_record(
    AGENT_ID,
    AGENT_NAME,
    ready=number(stats.get("totalPrompts")) > 0 or bool(limits["limits"]) or limits["signedIn"],
    # A sign-in with no sessions here has no local stats to show.
    hasLocalStats=number(stats.get("totalSessions")) > 0 or number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")),
    tierLabel=limits["tierLabel"],
    usageStatusText=limits["usageStatusText"],
    authHelpText=limits["authHelpText"],
    limits=limits["limits"],
    limitsStale=limits["stale"],
    limitsFetchedAt=limits["fetchedAtMs"],
  )
  record.update(stats)
  if limits["retryAdvised"]:
    record["retryAdvised"] = True
  emit(record)


if __name__ == "__main__":
  main()
