#!/usr/bin/env python3
"""claw-spend-guard v0.1 — hard daily/monthly spend cap for an OpenClaw gateway.

OpenClaw reports cost (openclaw status --usage, Control UI, usage-footer.json) but,
per docs read 2026-09-04 (concepts/usage-tracking, reference/api-usage-costs,
gateway/configuration), does not ENFORCE a dollar cap. This script does.

Cost sources, tried in order (each labelled with verification status):
  1. Anthropic Admin Cost API (ANTHROPIC_ADMIN_KEY)  — VERIFIED format (same API my
     own harness budget guard uses): GET /v1/organizations/cost_report.
  2. `openclaw status --usage --json`               — command VERIFIED in docs;
     JSON shape UNVERIFIED until tested on a local install (parsed defensively).
  3. ~/.openclaw/usage-footer.json                  — path VERIFIED in docs; shape UNVERIFIED.
  0. --cost-file PATH  (highest priority when given) — a JSON file {"day_usd": x, "month_usd": y}
     written by any exporter you trust (OpenRouter, your own logs). TESTED end-to-end.
Enforcement: stop the gateway unit (default `openclaw-gateway`, user systemd),
write a state file, send a Telegram alert; auto-resume at the cap's reset boundary.
Run from a systemd timer every 5 min. `--dry-run` reports without stopping anything.
"""
import argparse, datetime as dt, json, os, subprocess, sys, urllib.request

STATE = os.path.expanduser("~/.openclaw/spend-guard-state.json")

def utc_now():
    return dt.datetime.now(dt.timezone.utc)

def anthropic_cost_usd(admin_key, since):
    # Verified endpoint: org cost report, bucketed daily, amounts in USD cents-ish decimal strings.
    url = ("https://api.anthropic.com/v1/organizations/cost_report?starting_at="
           + since.strftime("%Y-%m-%dT00:00:00Z") + "&bucket_width=1d&limit=31")
    req = urllib.request.Request(url, headers={"x-api-key": admin_key, "anthropic-version": "2023-06-01"})
    with urllib.request.urlopen(req, timeout=20) as r:
        data = json.load(r)
    total = 0.0
    for bucket in data.get("data", []):
        for row in bucket.get("results", []):
            total += float(row.get("amount", 0)) / 100.0  # amount is in cents (decimal string)
    return total

def openclaw_status_cost():
    # UNVERIFIED shape: walk any JSON and sum keys that look like a USD cost total.
    try:
        out = subprocess.run(["openclaw", "status", "--usage", "--json"], capture_output=True, text=True, timeout=30).stdout
        data = json.loads(out)
    except Exception:
        return None
    found = []
    def walk(x):
        if isinstance(x, dict):
            for k, v in x.items():
                if k.lower() in ("cost_usd", "costusd", "total_cost", "totalcost", "spend") and isinstance(v, (int, float)):
                    found.append(float(v))
                walk(v)
        elif isinstance(x, list):
            for i in x: walk(i)
    walk(data)
    return max(found) if found else None

def telegram(msg):
    tok, chat = os.environ.get("TELEGRAM_BOT_TOKEN"), os.environ.get("TELEGRAM_CHAT_ID")
    if not tok or not chat: return
    body = json.dumps({"chat_id": chat, "text": msg}).encode()
    req = urllib.request.Request(f"https://api.telegram.org/bot{tok}/sendMessage", data=body, headers={"Content-Type": "application/json"})
    try: urllib.request.urlopen(req, timeout=10)
    except Exception: pass

def unit_cmd(action, unit, user_scope):
    cmd = ["systemctl"] + (["--user"] if user_scope else []) + [action, unit]
    return subprocess.run(cmd, capture_output=True, text=True).returncode == 0

def load_state():
    try: return json.load(open(STATE))
    except Exception: return {}

def save_state(s):
    d = os.path.dirname(os.path.abspath(STATE))
    os.makedirs(d, exist_ok=True)
    tmp = STATE + ".tmp"; json.dump(s, open(tmp, "w")); os.replace(tmp, STATE)  # atomic; a half-written state must never read as 'not stopped'

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--daily", type=float, required=True, help="daily hard cap in USD")
    ap.add_argument("--monthly", type=float, default=None, help="monthly hard cap in USD")
    ap.add_argument("--unit", default="openclaw-gateway", help="systemd unit name (verify with `systemctl --user list-units | grep -i openclaw`)")
    ap.add_argument("--system", action="store_true", help="unit is a system service, not --user")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--cost-file", default=None, help='JSON {"day_usd": x, "month_usd": y} from your own exporter')
    ap.add_argument("--state", default=None, help="state file path (default ~/.openclaw/spend-guard-state.json)")
    a = ap.parse_args()

    global STATE
    if a.state: STATE = a.state
    now = utc_now()
    src, day, month = None, None, None
    if a.cost_file:
        try:
            cf = json.load(open(a.cost_file))
            day, month, src = float(cf["day_usd"]), float(cf.get("month_usd", cf["day_usd"])), "cost-file"
        except Exception as e:
            print(f"cost-file unreadable: {e}", file=sys.stderr); sys.exit(2)
    key = os.environ.get("ANTHROPIC_ADMIN_KEY") or os.environ.get("ANTHROPIC_ADMIN_API_KEY")
    if key and day is None:
        try:
            day = anthropic_cost_usd(key, now)
            month = anthropic_cost_usd(key, now.replace(day=1))
            src = "anthropic-admin"
        except Exception as e:
            print(f"admin api failed: {e}", file=sys.stderr)
    if day is None:
        c = openclaw_status_cost()
        if c is not None:
            day, month, src = c, c, "openclaw-status(UNVERIFIED shape)"
    if day is None:
        print("no cost source available; refusing to guess", file=sys.stderr); sys.exit(2)

    state = load_state()
    over = (day >= a.daily) or (a.monthly is not None and month >= a.monthly)
    stopped = state.get("stopped_at")
    print(json.dumps({"ts": now.isoformat(), "source": src, "day_usd": round(day, 4), "month_usd": round(month or 0, 4),
                      "daily_cap": a.daily, "monthly_cap": a.monthly, "over": over, "stopped": bool(stopped), "dry_run": a.dry_run}))

    if over and not stopped:
        msg = f"claw-spend-guard: cap hit (day ${day:.2f}/{a.daily}, month ${month or 0:.2f}/{a.monthly}). Stopping {a.unit}."
        if not a.dry_run:
            state["stopped_at"] = now.isoformat(); state["stopped_day"] = now.strftime("%Y-%m-%d"); save_state(state)
            ok = unit_cmd("stop", a.unit, not a.system)
            if not ok: print(f"systemctl stop {a.unit} failed (check --system/--unit)", file=sys.stderr)
        telegram(msg + (" [dry-run]" if a.dry_run else ""))
    elif stopped and not over and state.get("stopped_day") != now.strftime("%Y-%m-%d"):
        # new UTC day and under cap → resume
        if not a.dry_run:
            unit_cmd("start", a.unit, not a.system)
            state.pop("stopped_at", None); save_state(state)
        telegram(f"claw-spend-guard: new day, under cap → resumed {a.unit}.")

if __name__ == "__main__":
    main()