# usage-statusline.py - Andre's "how much have I got left" tool, shared so you can have it too.
#
# GIVE THIS WHOLE FILE TO YOUR ASSISTANT and say: "adapt this to my setup so my Claude Code
# status line shows my model, my effort level, and how long until my token window resets."
# It needs two things wired to YOUR account: your claude.ai org id (REPLACE_WITH_YOUR_ORG_ID)
# and a way to read your own claude.ai session (REPLACE_WITH_YOUR_BROWSER_PROFILE). Your
# assistant figures that out for your machine. Nothing here is a secret - it only reads your
# own usage. The mode you want is:  python3 usage-statusline.py --statusline
#
#!/usr/bin/env python3
"""
claude-usage — your token-headroom gauge.

Reads your REAL Claude subscription usage from claude.ai's (unofficial) usage
endpoint and reports headroom, so you can see how much headroom you've got.


Architecture: ONE authed fetcher (root, reads your browser's claude.ai cookie
path) refreshes a small secret-free JSON cache; MANY cheap readers (login
greeting, glut cron, `claude-usage` CLI) read the cache — no network, no creds.

  claude-usage              one-line gauge for humans (reads cache)
  claude-usage --json       machine output for crons (reads cache)
  claude-usage --fetch      refresh the cache from the endpoint (needs root)

Endpoint (confirmed S118): GET https://claude.ai/api/organizations/{org}/usage
  -> {"five_hour": {"utilization": 25.0, "resets_at": "..."}, "seven_day": {...},
      "seven_day_opus": null, "seven_day_sonnet": {...}, ...}
Auth: full claude.ai cookie jar from root's Firefox profile (cf_clearance etc.),
(your assistant wires this to your own claude.ai session).
"""
import sqlite3, shutil, tempfile, os, json, sys, time, datetime, urllib.request, logging

PROFILE   = os.environ.get("CLAUDE_USAGE_PROFILE", "REPLACE_WITH_YOUR_BROWSER_PROFILE")
ORG       = os.environ.get("CLAUDE_USAGE_ORG", "REPLACE_WITH_YOUR_ORG_ID")
UA        = "Mozilla/5.0 (X11; Linux x86_64; rv:140.0) Gecko/20100101 Firefox/140.0"
BASE      = "https://claude.ai"
STATE     = os.path.expanduser("~/.claude-usage/state.json")   # secret-free cache, world-readable
LOG       = os.path.expanduser("~/.claude-usage/usage.log")
STALE_MIN = 45        # cache older than this => degraded
QUIET_PCT = float(os.environ.get("CLAUDE_USAGE_QUIET_PCT", "40"))  # both windows below = quiet (tunable)

logger = logging.getLogger("claude-usage")
logger.setLevel(logging.INFO)
def _setup_log():
    """Attach a file handler only if the log is writable (fetch runs as root and
    can; cheap readers as brodawg can't touch the root-owned log — they no-op)."""
    if logger.handlers:
        return
    try:
        h = logging.FileHandler(LOG)
        h.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
        logger.addHandler(h)
    except Exception:
        logger.addHandler(logging.NullHandler())


# ---------- fetch (root only) ----------
def _cookies():
    with tempfile.TemporaryDirectory() as t:
        dst = os.path.join(t, "c.sqlite")
        shutil.copy(os.path.join(PROFILE, "cookies.sqlite"), dst)
        con = sqlite3.connect(dst)
        rows = con.execute("SELECT name,value FROM moz_cookies WHERE host LIKE '%claude.ai%'").fetchall()
        con.close()
    return "; ".join(f"{n}={v}" for n, v in rows)


def fetch():
    """Hit the endpoint, write the secret-free cache. Returns the windows dict."""
    hdr = {"User-Agent": UA, "Cookie": _cookies(), "Accept": "*/*",
           "Referer": f"{BASE}/", "Origin": BASE, "Sec-Fetch-Site": "same-origin",
           "Sec-Fetch-Mode": "cors", "Sec-Fetch-Dest": "empty"}
    url = f"{BASE}/api/organizations/{ORG}/usage"
    r = urllib.request.Request(url, headers=hdr)
    with urllib.request.urlopen(r, timeout=40) as resp:
        raw = json.load(resp)
    # keep ONLY the window utilisation + reset fields — never any cookie/secret
    keep = {}
    for k in ("five_hour", "seven_day", "seven_day_opus", "seven_day_sonnet"):
        w = raw.get(k)
        keep[k] = ({"utilization": w.get("utilization"), "resets_at": w.get("resets_at")}
                   if isinstance(w, dict) else None)
    out = {"windows": keep, "fetched_at": time.time()}
    tmp = STATE + ".tmp"
    with open(tmp, "w") as f:
        json.dump(out, f)
    os.replace(tmp, STATE)
    os.chmod(STATE, 0o644)
    logger.info("fetched ok five_hour=%s seven_day=%s",
                 keep["five_hour"] and keep["five_hour"]["utilization"],
                 keep["seven_day"] and keep["seven_day"]["utilization"])
    return out


# ---------- read (anyone) ----------
def load():
    """Return (state_dict, degraded_reason|None)."""
    if not os.path.exists(STATE):
        return None, "no cache yet (run --fetch)"
    try:
        with open(STATE) as f:
            st = json.load(f)
    except Exception as e:
        return None, f"cache unreadable: {e}"
    age_min = (time.time() - st.get("fetched_at", 0)) / 60.0
    if age_min > STALE_MIN:
        return st, f"stale ({age_min:.0f} min old)"
    return st, None


def _util(st, key):
    w = (st.get("windows") or {}).get(key)
    return (w or {}).get("utilization")


def _countdown(iso):
    if not iso:
        return "—"
    try:
        t = datetime.datetime.fromisoformat(iso.replace("Z", "+00:00"))
        secs = (t - datetime.datetime.now(datetime.timezone.utc)).total_seconds()
        if secs <= 0:
            return "due"
        h = int(secs // 3600); m = int((secs % 3600) // 60)
        return f"{h}h{m:02d}m" if h else f"{m}m"
    except Exception:
        return "?"


def _clock(iso):
    """Reset timestamp -> local 24hr wall-clock, e.g. '14:35'. Shown beside the
    countdown so a stalled session can glance at the actual turnover time."""
    if not iso:
        return None
    try:
        t = datetime.datetime.fromisoformat(iso.replace("Z", "+00:00"))
        return t.astimezone().strftime("%H:%M")
    except Exception:
        return None


def _week_elapsed(iso):
    """How far through the current weekly window we are, as a %. Window is 7 days
    ending at resets_at, so start = resets_at − 7d. Lets '3% used / 40% elapsed'
    show pace — the bigger (elapsed − used), the more weekly slack."""
    if not iso:
        return None
    try:
        reset = datetime.datetime.fromisoformat(iso.replace("Z", "+00:00"))
        start = reset - datetime.timedelta(days=7)
        now = datetime.datetime.now(datetime.timezone.utc)
        return max(0.0, min(100.0, (now - start).total_seconds() / (7 * 86400) * 100))
    except Exception:
        return None


def _run_dry(used, elapsed, reset_iso):
    """Time-aware weekly read: project when the weekly tokens run dry at the
    current within-window average rate, and compare to when the window resets.
    Replaces the consequence-blind pace-gap — the SAME used-vs-elapsed gap reads
    as a bigger 'before reset' figure late in the week (less runway to recover)
    and a small one early, emergent from the projection (no day-of-week logic).

    Model: used% grows ~linearly from window start, so cumulative hits 100% at
    exhaust_frac = elapsed%/used% of the 7-day window. <1 → dries early (⚠);
    ≥1 (incl. nothing-used / ahead-of-pace) → coasts ✅. Aggregate-only: no burn
    history needed. Returns "" only when the reset timestamp is unparseable."""
    if not reset_iso or used is None or elapsed is None or used <= 0:
        return "coasts ✅"
    try:
        reset = datetime.datetime.fromisoformat(reset_iso.replace("Z", "+00:00"))
    except Exception:
        return ""
    start = reset - datetime.timedelta(days=7)
    exhaust_frac = elapsed / used            # fraction of the 7-day window
    if exhaust_frac >= 1.0:
        return "coasts ✅"
    dry = start + datetime.timedelta(days=7 * exhaust_frac)
    secs = (reset - dry).total_seconds()     # always > 0 here
    delta = (f"~{secs / 86400:.0f}d before reset" if secs >= 72 * 3600
             else f"~{secs / 3600:.0f}h before reset")
    when = dry.astimezone().strftime("%a %H:%M")   # local wall-clock, e.g. 'Fri 15:00'
    return f"dries {when} ({delta} ⚠)"


def _allowance(used, elapsed):
    """The weekly read that matches Andre's 'late is worse' intuition (S129): the
    fraction of NORMAL pace you can still spend and land on budget by reset.
        ratio   = (100 − used%) / (100 − elapsed%)   — remaining budget per remaining time
        per_day = (100 − used%) / remaining_days       — concrete %/day still allowed
    Unlike the burn-rate-driven run-dry clock, this tightens monotonically as the
    week runs out for a fixed used-vs-elapsed gap (more calendar runway early = less
    tragic), which is the consequence-weighting Andre asked for. Aggregate-only —
    no burn history. Returns (ratio, per_day_pct, label). ratio ≥ 1 → room to spare.
    The ratio is the number a future glutifier/throttler can read straight from --json."""
    if used is None or elapsed is None or elapsed >= 100:
        return None, None, ""
    rem_budget = max(0.0, 100.0 - used)
    rem_days = 7.0 * (100.0 - elapsed) / 100.0
    per_day = (rem_budget / rem_days) if rem_days > 0 else None
    ratio = rem_budget / (100.0 - elapsed)
    if used >= 100:
        label = "🔴 weekly cap hit — wait for reset"
    elif ratio >= 3.0:
        label = "✅ room to spare"
    elif ratio >= 1.05:
        label = f"✅ {ratio:.1f}× normal pace"
    elif ratio >= 1.0:
        label = "✅ on pace"
    else:
        label = f"ease to {ratio * 100:.0f}% of normal ⚠"
    return round(ratio, 2), (round(per_day, 1) if per_day is not None else None), label


def _tier(binding):
    """Plain-language 'how much can I spend right now', by the binding (busiest)
    window's utilisation. This is the judgement any instance reads to match a task
    to current headroom — replaces the old jargon 'glut' bool."""
    if binding < 40:
        return "big", "🟢 room for a big one"
    if binding < 80:
        return "small", "🟡 small tasks only"
    return "lifelines", "🔴 lifelines only"


def gauge(st):
    """Compute headroom + quiet flag from a loaded state dict."""
    fh = _util(st, "five_hour") or 0.0
    sd = _util(st, "seven_day") or 0.0
    binding = max(fh, sd)
    headroom = round(100 - binding, 1)
    quiet = (fh < QUIET_PCT and sd < QUIET_PCT)
    fh_w = (st["windows"].get("five_hour") or {})
    sd_w = (st["windows"].get("seven_day") or {})
    sd_elapsed = _week_elapsed(sd_w.get("resets_at"))
    _allowance_ratio, _allowance_per_day, _allowance_label = _allowance(sd, sd_elapsed)
    tier, tier_label = _tier(binding)
    return {
        "five_hour_pct": fh, "seven_day_pct": sd,
        "seven_day_elapsed_pct": sd_elapsed,
        # seven_day_slack (elapsed% − used%) stays as the internal number a future
        # glutifier reads — slack ≥ 20 is its weekly fire threshold. Unchanged.
        "seven_day_slack": (round(sd_elapsed - sd, 1) if sd_elapsed is not None else None),
        # The weekly severity read (S129): allowance leads the human bar AND is the
        # machine-checkable number in --json (seven_day_pace_ratio). run_dry kept as a
        # secondary projection. See _allowance / _run_dry for the why.
        "seven_day_pace_ratio": _allowance_ratio,
        "seven_day_allowance_per_day": _allowance_per_day,
        "seven_day_pace_label": _allowance_label,
        "seven_day_run_dry": _run_dry(sd, sd_elapsed, sd_w.get("resets_at")),
        "headroom_pct": headroom, "glut": quiet,
        "tier": tier, "tier_label": tier_label,
        "five_hour_resets_in": _countdown(fh_w.get("resets_at")),
        "five_hour_reset_clock": _clock(fh_w.get("resets_at")),
        "seven_day_resets_in": _countdown(sd_w.get("resets_at")),
    }


def human_line(st, degraded):
    g = gauge(st)
    flag = " ⚠degraded" if degraded else ""
    wk = f"week {g['seven_day_pct']:.0f}% used"
    if g["seven_day_pace_label"]:
        wk += f" · {g['seven_day_pace_label']}"
    if g["seven_day_run_dry"]:
        wk += f" · {g['seven_day_run_dry']}"
    return (f"usage · 5h {g['five_hour_pct']:.0f}% used · {g['five_hour_resets_in']} to reset · "
            f"{wk}, resets {g['seven_day_resets_in']} · "
            f"headroom {g['headroom_pct']:.0f}% · {g['tier_label']}{flag}")


# ---------- glut-by-absence fallback (no endpoint) ----------
def fallback_line():
    """If the endpoint/cache is unavailable: estimate from local transcript token
    sums over the trailing 5h. Honest, degraded — not the real %."""
    import glob
    d = os.path.expanduser("~/.claude/projects/-home-brodawg")
    cutoff = time.time() - 5 * 3600
    toks = 0
    for fp in glob.glob(os.path.join(d, "*.jsonl")):
        if os.path.getmtime(fp) < cutoff:
            continue
        try:
            with open(fp, encoding="utf-8", errors="replace") as f:
                for line in f:
                    if '"usage"' not in line:
                        continue
                    try:
                        u = (json.loads(line).get("message", {}) or {}).get("usage") or {}
                        toks += (u.get("input_tokens", 0) or 0) + (u.get("output_tokens", 0) or 0)
                    except Exception:
                        pass
        except OSError:
            pass
    return f"usage · ENDPOINT DOWN — glut-by-absence: ~{toks:,} tokens in last 5h (no % available)"


def main(argv):
    _setup_log()
    if "--fetch" in argv:
        try:
            fetch()
            if "--quiet" not in argv:
                st, deg = load(); print(human_line(st, deg))
            return 0
        except Exception as e:
            logger.error("fetch failed: %s", e)
            print(f"claude-usage: fetch failed: {e}", file=sys.stderr)
            return 1
    st, deg = load()
    if "--statusline" in argv:
        # Claude Code statusline: reads its context JSON on stdin (model, dir),
        # appends the usage gauge. MUST NOT error or the statusline breaks — every
        # step guarded, degrades to whatever it can show.
        model = effort = cwd = ""
        try:
            ctx = json.load(sys.stdin)
            model = (ctx.get("model") or {}).get("display_name") or ""
            # effort.level: low/medium/high/xhigh/max; object absent when the
            # model doesn't support effort, so guard and skip silently if missing.
            effort = (ctx.get("effort") or {}).get("level") or ""
            ws = ctx.get("workspace") or {}
            cwd = ws.get("current_dir") or ctx.get("cwd") or ""
        except Exception:
            pass
        parts = ["assistant"]
        if model:
            parts.append(model)
        if effort:
            parts.append(effort)
        if cwd:
            parts.append(os.path.basename(cwd.rstrip("/")) or cwd)
        if st is not None:
            g = gauge(st)
            flag = " ⚠" if deg else ""
            lbl = g["seven_day_pace_label"]
            wk = f"wk {g['seven_day_pct']:.0f}% · {lbl}" if lbl else f"wk {g['seven_day_pct']:.0f}%"
            clk = g["five_hour_reset_clock"]
            rst = f"{g['five_hour_resets_in']} left" + (f" ({clk})" if clk else "")
            parts.append(f"5h {g['five_hour_pct']:.0f}% · {rst} · {wk} · {g['tier_label']}{flag}")
        else:
            parts.append("usage —")
        print(" · ".join(parts))
        return 0
    if "--greet" in argv:
        # compact, fast, silent-on-no-cache — for the login greeting. Never noisy.
        if st is None:
            return 0
        g = gauge(st)
        flag = " ⚠" if deg else ""
        lbl = g["seven_day_pace_label"]
        wk = f"week {g['seven_day_pct']:.0f}% · {lbl}" if lbl else f"week {g['seven_day_pct']:.0f}%"
        clk = g["five_hour_reset_clock"]
        rst = f"{g['five_hour_resets_in']} left" + (f" ({clk})" if clk else "")
        print(f"usage · 5h {g['five_hour_pct']:.0f}% · {rst} · {wk} · {g['tier_label']}{flag}")
        return 0
    if st is None:
        # no usable cache — degrade honestly
        if "--json" in argv:
            print(json.dumps({"error": deg, "fallback": fallback_line()})); return 1
        print(fallback_line()); return 1
    if "--json" in argv:
        out = gauge(st); out["degraded"] = deg; print(json.dumps(out)); return 0
    print(human_line(st, deg))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))