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

Everything the agents panel shows for Ollama Cloud comes from this one
command: local token stats from pi and omp sessions, opencode's message
database, Codex's `--oss` rollouts and Claude Code's transcripts that ran on
the `ollama-cloud` provider (or on a cloud model through the local daemon),
and the account's 5-hour session and 7-day weekly
limits (or the monthly included credits) from Ollama's balance endpoint. The
panel itself only ever reads the JSON this prints; it never talks to disk
formats or endpoints.
"""

from __future__ import annotations

import argparse
import math
import os
import sys
import urllib.request
from datetime import datetime
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,
  base_record,
  cached_scan as shared_cached_scan,
  cloud_model,
  config_home,
  emit,
  limits_cache_path as shared_limits_cache_path,
  open_json,
  opencode_db,
  pi_session_roots,
  read_json,
  routed_match,
  routed_sources,
  scan_max_age,
  scan_opencode,
  scan_pi_sessions,
  scan_routed,
)

AGENT_ID = "ollama"
AGENT_NAME = "Ollama Cloud"
PROVIDER = "ollama-cloud"
BALANCE_ENDPOINT = "https://ollama.com/api/balance"
ME_ENDPOINT = "https://ollama.com/api/me"
AUTH_HELP_NO_KEY = "Set OLLAMA_API_KEY, or sign in to pi with an Ollama Cloud key, to show the account limits."
AUTH_HELP_BAD_KEY = (
  "Ollama Cloud rejected the API key. Update OLLAMA_API_KEY or ~/.config/omarchy/agents/ollama.json,"
  " or sign in to pi again. Local stats are still shown."
)
# GET /api/balance (docs.ollama.com/api/balance) answers a legacy plan with
# `included.session` / `included.weekly` ({remaining_percent 0..100,
# resets_at}) and a current plan with `included.balance_usd`,
# `allowance_usd` and `period.until`. /api/usage only reports token and cost
# history, never the limits.


def config_dir() -> Path:
  return config_home() / "omarchy" / "agents"


def read_config() -> dict[str, Any]:
  return read_json(config_dir() / "ollama.json")


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


def api_key() -> str:
  env = os.environ.get("OLLAMA_API_KEY", "").strip()
  if env:
    return env
  configured = str(read_config().get("apiKey") or "").strip()
  if configured:
    return configured
  # pi and omp store the Ollama Cloud key they signed in with; either one
  # proves the account, and the balance endpoint accepts the same key. Only a
  # literal key is taken: a `$VAR` or `!command` entry is resolved by pi alone.
  for auth_path in (pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"):
    entry = read_json(auth_path).get(PROVIDER)
    key = str(entry.get("key") or "").strip() if isinstance(entry, dict) else ""
    if key and not key.startswith(("!", "$")):
      return key
  return ""


# ---------------------------------------------------------------- local scan
#
# pi and omp record every assistant message with its provider, model, and
# token usage in JSONL session files, and opencode in its message database,
# under the same provider ids. The `ollama-cloud` provider is the
# Ollama Cloud API, which serves Claude models alongside the open ones. A
# cloud-tagged model ("glm-4.6:cloud", "gpt-oss:120b-cloud") reached through
# the local daemon's `ollama` provider runs on Ollama's servers too, so it
# counts here; the local collector leaves exactly those out. Codex's `--oss`
# rollouts name the same `ollama` provider, and count here on the same
# models. Claude Code pointed at ollama.com, or answering with a cloud-tagged
# model, is Ollama's too, as is a MiniMax Code custom provider at ollama.com.


def billed_by_ollama_cloud(provider: str, model: str, entry: Any = None) -> bool:
  return provider == PROVIDER or (provider == "ollama" and cloud_model(model))


# Ollama model ids carry a version tag ("deepseek-v4-pro:0813"); merge tags
# so the chart groups by model, not by release.
def model_name(raw: str) -> str:
  return (raw or "ollama").split(":", 1)[0]


def scan_sessions(model_prefix: str, reuse: bool = True) -> tuple[dict[str, Any], bool]:
  stats = Stats()
  pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"ollama(-cloud)?"', billed_by_ollama_cloud)
  # The database's prefilter matches both provider ids.
  db_rows, db_complete = scan_opencode(["ollama"], billed_by_ollama_cloud)
  # Codex rollouts come with them, by codex_route_owner(): the `--oss`
  # daemon's cloud-tagged models, a model_providers entry named
  # `ollama-cloud`, or one whose base_url is ollama.com.
  routed_rows, routed_complete = scan_routed(AGENT_ID, reuse)
  for row in pi_rows + db_rows + routed_rows:
    model = model_name(row.model)
    if model_prefix and not model.lower().startswith(model_prefix):
      continue
    stats.add_row(row, model)
  return stats.snapshot(), pi_complete and db_complete and routed_complete


def cached_scan(max_age_seconds: float, model_prefix: str) -> dict[str, Any]:
  # The digest covers every data path the scan reads: the pi/omp session
  # roots, the opencode DB, Codex's home, and Claude Code's and MiniMax
  # Code's sessions. A scan only answers for the model filter it applied, and
  # the endpoints those were sent to.
  return shared_cached_scan(
    AGENT_ID,
    [*pi_session_roots(), opencode_db(), *routed_sources()],
    max_age_seconds,
    lambda: scan_sessions(model_prefix, reuse=max_age_seconds > 0),
    match={"modelPrefix": model_prefix, **routed_match()},
  )


# ------------------------------------------------------------------- limits


def clamp_fraction(value: Any) -> float:
  if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0:
    return -1.0
  return min(1.0, float(value))


def reset_time(value: Any) -> str:
  raw = value.strip() if isinstance(value, str) else ""
  try:
    datetime.fromisoformat(raw.replace("Z", "+00:00"))
  except ValueError:
    return ""
  return raw


def amount(value: Any) -> float | None:
  if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
    return None
  return float(value)


def balance_limits(payload: Any) -> list[dict[str, Any]]:
  included = payload.get("included") if isinstance(payload, dict) else None
  if not isinstance(included, dict):
    return []
  out = []
  for key, label in (("session", "Session (5-hour)"), ("weekly", "Weekly (7-day)")):
    window = included.get(key)
    if not isinstance(window, dict):
      continue
    # remaining_percent runs 0..100 and says what is left, not what is used.
    remaining = amount(window.get("remaining_percent"))
    remaining = clamp_fraction(remaining / 100) if remaining is not None else -1.0
    if remaining >= 0:
      out.append({"label": label, "percent": 1.0 - remaining, "resetsAt": reset_time(window.get("resets_at"))})
  balance, allowance = amount(included.get("balance_usd")), amount(included.get("allowance_usd"))
  if balance is not None and allowance is not None and allowance > 0:
    period = included.get("period")
    out.append({
      "label": "Included credits",
      "title": "Monthly",
      "percent": max(0.0, min(1.0, (allowance - balance) / allowance)),
      "resetsAt": reset_time(period.get("until")) if isinstance(period, dict) else "",
    })
  return out


# The credits bought apart from the plan: purchased.balance_usd, what's left
# of them unexpired. The answer says how much remains but not how much was
# bought, so the balance shows the remaining figure alone. A plan that never
# bought any, or used them up, has none to show, and that is no alarm while
# the plan's included credits still serve.
def purchased_balance(payload: Any) -> dict[str, Any] | None:
  purchased = payload.get("purchased") if isinstance(payload, dict) else None
  remaining = amount(purchased.get("balance_usd")) if isinstance(purchased, dict) else None
  if remaining is None or remaining <= 0:
    return None
  return {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False}


def kept_balance(cache: LimitsCache) -> dict[str, Any] | None:
  balance = cache.data.get("balance")
  return balance if isinstance(balance, dict) else None


def probe_limits(access_key: str) -> dict[str, Any]:
  request = urllib.request.Request(
    BALANCE_ENDPOINT,
    headers={"Authorization": "Bearer " + access_key, "Accept": "application/json"},
  )
  try:
    payload = open_json(request, timeout=10)
  except HttpError as error:
    if error.status in (401, 403):
      return {"ok": False, "auth": True, "helpText": AUTH_HELP_BAD_KEY}
    if error.status:
      return {"ok": False, "helpText": f"Ollama's balance endpoint returned status {error.status}. Local stats are still shown."}
    if error.transport:
      # A transport failure reached no server at all — no route, no DNS. Any
      # real answer, including an error status, is a server we should stop
      # pestering; this is not.
      return {
        "ok": False,
        "transport": True,
        "helpText": "Couldn't reach Ollama's balance endpoint. Retrying shortly. Local stats are still shown.",
      }
    return {"ok": False, "helpText": "Ollama's balance endpoint returned an invalid answer. Local stats are still shown."}

  out = balance_limits(payload)
  balance = purchased_balance(payload)
  if not out and balance is None:
    return {"ok": False, "helpText": "Ollama's balance endpoint returned no limits. Local stats are still shown."}
  return {"ok": True, "limits": out, "balance": balance}


def probe_plan(access_key: str) -> str | None:
  """The plan Ollama names for the key, "" when it names none, None when it couldn't be asked."""
  request = urllib.request.Request(
    ME_ENDPOINT,
    headers={"Authorization": "Bearer " + access_key, "Accept": "application/json"},
    method="POST",
  )
  try:
    payload = open_json(request, timeout=10)
  except Exception:
    return None
  # Ollama's UserResponse tags the field `plan`; Go's decoder takes either
  # case, so `Plan` is read too.
  plan = str(payload.get("plan") or payload.get("Plan") or "").strip()[:64] if isinstance(payload, dict) else ""
  return plan[0].upper() + plan[1:] if plan else ""


def limits_cache_path(access_key: str) -> Path:
  # Keyed by the credential: another key never shows this one's allowance or
  # plan.
  return shared_limits_cache_path(AGENT_ID, access_key)


def collect_limits(access_key: str, force: bool) -> dict[str, Any]:
  # `live` says whether the limits are this check's (or a probe recent enough
  # to stand in for one), `fetchedAtMs` when they were measured, so the panel
  # can tell kept limits from fresh ones.
  result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "fetchedAtMs": 0}

  if access_key == "":
    result["usageStatusText"] = "Waiting for auth"
    result["authHelpText"] = AUTH_HELP_NO_KEY
    return result

  # A panel that is opened and shut repeatedly must not turn into a request
  # per flick, so recent probe results are reused for a short window — and
  # kept as the answer of record when a later probe fails.
  cache = LimitsCache(AGENT_ID, access_key)
  cached_plan = str(cache.data.get("plan") or "")

  # --force is a person asking for fresh numbers, so it skips the reuse window
  # entirely; the interval is there to absorb repeated panel opens, not to
  # overrule someone who pressed refresh.
  if cache.fresh(force, allow_empty=kept_balance(cache) is not None):
    result.update(limits=cache.limits, tierLabel=cached_plan, live=True, fetchedAtMs=cache.fetched_at, balance=kept_balance(cache))
    return result

  probe = probe_limits(access_key)
  if probe["ok"]:
    plan = probe_plan(access_key)
    # A plan lookup that failed keeps the plan last named, rather than
    # blanking a label the account still has.
    plan = cached_plan if plan is None else plan
    fetched_at = cache.save(probe["limits"], only_if_newer=True, plan=plan, balance=probe["balance"])
    result.update(limits=probe["limits"], tierLabel=plan, live=True, fetchedAtMs=fetched_at, balance=probe["balance"])
    return result

  # The first probe after login often fires before DHCP has handed out a
  # route. Ask the shell to try again sooner than its regular interval.
  if probe.get("transport"):
    result["retryAdvised"] = True
  result["authHelpText"] = probe["helpText"]
  fallback = cache.kept()[0]
  balance = kept_balance(cache)
  if fallback or balance:
    result.update(limits=fallback, tierLabel=cached_plan, fetchedAtMs=cache.fetched_at, balance=balance)
  if probe.get("auth"):
    result["usageStatusText"] = "Sign-in expired"
  elif not fallback and not balance:
    result["usageStatusText"] = "Ollama limits unavailable"
  return result


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


def main() -> int:
  parser = argparse.ArgumentParser()
  parser.add_argument("--force", action="store_true", help="rescan sessions and re-probe limits, ignoring caches")
  parser.add_argument("--limits-only", action="store_true", help="reuse any recent session scan; only the limits probe must be fresh")
  args = parser.parse_args()

  config = read_config()
  model_prefix = str(config.get("modelPrefix") or "").strip().lower()
  stats = cached_scan(scan_max_age(args.force, args.limits_only), model_prefix)
  has_stats = bool(stats.get("modelUsage"))

  key = api_key()
  record = base_record(AGENT_ID, AGENT_NAME, hasLocalStats=has_stats)
  record.update(stats)

  # A machine with neither a key nor any Ollama Cloud sessions doesn't use it,
  # and the panel skips the record; that is no trouble to report.
  if key or has_stats:
    limits = collect_limits(key, args.force)
    record.update(
      ready=True,
      tierLabel=limits["tierLabel"],
      usageStatusText=limits["usageStatusText"],
      authHelpText=limits["authHelpText"],
      limits=limits["limits"],
      limitsStale=bool(limits["limits"] or limits.get("balance")) and not limits["live"],
      limitsFetchedAt=limits["fetchedAtMs"],
    )
    if limits.get("balance"):
      record["balance"] = limits["balance"]
    if limits.get("retryAdvised"):
      record["retryAdvised"] = True
  emit(record, sort_keys=True)
  return 0


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