ocdway/scripts/ocd-waybar

#!/usr/bin/env python3
# ocd-waybar: waybar daemon for the ocd opencode-limits dashboard.
#
# It polls `ocd raw json` every 30 s. It does not pass --force-refresh or
# --no-cache, so ocd's own cache and TTL stay in charge of politeness.
# It holds the 5 s rotation timer and the 1 h lock, renders one Pango line
# for the displayed account and a multi-line Pango tooltip for all accounts,
# and writes the result as JSON for the waybar custom/ocd module.
#
# Clicks arrive over a unix socket from ocd-waybar-ctl:
#   advance  -> reset the 5 s timer and move to the next account
#   lock     -> freeze the current account for one hour
#   poll     -> force an ocd poll now (still respects ocd's cache/TTL)
#   quit     -> stop the daemon

import json
import os
import re
import signal
import socket
import subprocess
import sys
import threading
import time
from datetime import datetime, timezone
from html import escape
from pathlib import Path

# --- tunables ---------------------------------------------------------------
ROTATE_SECS   = 5        # normal auto-advance cadence
LOCK_SECS     = 3600     # one hour lock
POLL_SECS     = 30       # ocd re-poll cadence (cache serves most calls)
TICK_SECS     = 0.5      # main loop resolution
BAR_WIDTH     = 7        # nbsp cells per bar (~50 px at normal bar font)
STATE_DIR     = Path.home() / ".local" / "state" / "ocd-waybar"
STATE_FILE    = STATE_DIR / "state.json"
SOCK_FILE     = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")) / "ocd-waybar.sock"
OCD_BIN        = os.environ.get("OCD_BIN", "ocd")
# Default False: waybar 0.15.0 renders the custom-module tooltip as plain
# text, so tags would show literally. Set True if your waybar supports
# Pango markup in custom tooltips (see Waybar PR #5034).
TOOLTIP_MARKUP = False

# zone colors
C_OK   = "#43c267"
C_WARN = "#e0a030"
C_CRIT = "#e0533b"
C_LOCK = "#5aa6e0"
C_TRACK = "#39424d"

ZONE = {("ok",): 0, ("warn",): 1, ("crit",): 2}
# Each label is a template: {} is filled with the used value plus its
# unit, e.g. 88 and "%" -> "[88%/5h]". With no data the unit alone shows.
FRAME_LABELS = (("usage_5h", "[{}/5h]"), ("usage_week", "[{}/1w]"), ("usage_month", "[{}/1m]"))


def zone_color(pct, status=None, locked=False):
    if locked:
        return C_LOCK
    if status == "rate-limited" or pct >= 90:
        return C_CRIT
    if pct >= 70:
        return C_WARN
    return C_OK


def zone_name(pct, status=None):
    if status == "rate-limited" or pct >= 90:
        return "ocd-crit"
    if pct >= 70:
        return "ocd-warn"
    return "ocd-ok"


def parse_dt(s):
    """Parse an ISO 8601 string with possible nanosecond precision."""
    if not s:
        return None
    s = s.strip()
    if s.endswith("Z"):
        s = s[:-1] + "+00:00"
    # Python only handles up to 6 fractional digits. Trim the rest.
    m = re.search(r"(\.\d{6})\d+", s)
    if m:
        s = s[:m.start()] + m.group(1) + s[m.end():]
    try:
        dt = datetime.fromisoformat(s)
    except ValueError:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt


def rel_reset(resets_at, now):
    """Human string for time until reset, e.g. '4:39', '3d 1h', '20d'."""
    if not resets_at:
        return "\u2014"
    dt = parse_dt(resets_at)
    if dt is None:
        return "\u2014"
    sec = (dt - now).total_seconds()
    if sec <= 0:
        return "now"
    days = int(sec // 86400)
    rem = sec - days * 86400
    hours = int(rem // 3600)
    mins = int((rem % 3600) // 60)
    if days >= 7:
        return f"{days}d"
    if days >= 1:
        h = hours if hours else 0
        return f"{days}d {h}h"
    return f"{hours}:{mins:02d}"


def rel_ago(fetched_at, now):
    if not fetched_at:
        return "\u2014"
    dt = parse_dt(fetched_at)
    if dt is None:
        return "\u2014"
    sec = (now - dt).total_seconds()
    if sec < 0:
        sec = 0
    if sec < 60:
        return f"{int(sec)}s ago"
    m = int(sec // 60)
    if m < 60:
        return f"{max(1, m)}m ago"
    h = int(sec // 3600)
    if h < 24:
        return f"{h}h ago"
    return f"{h // 24}d ago"


def bar_html(pct, status=None, locked=False):
    pct = max(0.0, min(100.0, float(pct)))
    fill = int(round(pct / 100.0 * BAR_WIDTH))
    fill = max(0, min(BAR_WIDTH, fill))
    empty = BAR_WIDTH - fill
    fc = zone_color(pct, status, locked)
    out = f'<span background="{fc}">{"\u00a0" * fill}</span>'
    if empty:
        out += f'<span background="{C_TRACK}">{"\u00a0" * empty}</span>'
    return out


def bar_plain(pct, status=None, locked=False):
    """Plain block-character bar for tooltip text that may not parse markup."""
    pct = max(0.0, min(100.0, float(pct)))
    units = pct / 100.0 * BAR_WIDTH            # cells, fractional
    full = int(units)
    frac = units - full
    blocks = "▏▎▍▌▋▊▉"  # one-eighth steps
    out = "█" * full
    if full < BAR_WIDTH and frac > 0:
        idx = int(round(frac * 7)) - 1     # 7 fractional steps
        idx = max(0, min(6, idx))
        out += blocks[idx]
        full += 1
    out += "░" * (BAR_WIDTH - full)    # light shade for the empty track
    if locked and full < BAR_WIDTH:
        out = out[:full] + "▓" + out[full + 1:]  # dim mark on locked rows
    return out[:BAR_WIDTH]


def frame_get(limits, key):
    if not limits:
        return None
    return limits.get(key)


def pct_of(frame):
    if frame is None:
        return 0.0
    total = float(frame.get("total") or 0) or 100.0
    used = float(frame.get("used") or 0)
    return used / total * 100.0 if total else 0.0


def worst_zone(account, locked=False):
    if not account or not account.get("limits"):
        return "ocd-ok"
    worst = 0
    names = ("usage_5h", "usage_week", "usage_month")
    for k in names:
        f = account["limits"].get(k)
        if not f:
            continue
        zn = zone_name(pct_of(f), f.get("status"))
        w = {"ocd-ok": 0, "ocd-warn": 1, "ocd-crit": 2}[zn]
        worst = max(worst, w)
    return ("ocd-ok", "ocd-warn", "ocd-crit")[worst]


def balance_str(account):
    bal = account.get("balance") if account else None
    if not bal:
        return "\u2014"
    remaining = float(bal.get("remaining") or 0.0)
    unit = bal.get("unit") or ""
    # Currency symbol prefixes; others append a space.
    if unit and not unit.isalpha():
        return f"{unit}{remaining:.2f}"
    return f"{remaining:.2f} {unit}".strip()


def render_account_line(account, now, locked, markup=True, with_name=True):
    """Render one account: name + lock emoji (optional), then the three
    bars with their percentage labels, then balance and update age."""
    limits = account.get("limits") if account else None
    parts = []
    if with_name:
        nm = account.get("account") if account else ""
        if nm:
            nm = escape(nm) if markup else nm
            parts.append(("\U0001f512 " if locked else "") + nm)
    segs = []
    for key, label in FRAME_LABELS:
        f = frame_get(limits, key)
        if f is None:
            filled = label.format("\u2014")
            bar = "" if markup else ""
            rel = "\u2014"
        else:
            used = f.get("used") or 0.0
            unit = f.get("unit") or ""
            val = f"{int(round(used))}{unit}"
            filled = label.format(val)
            bar = bar_html(pct_of(f), f.get("status"), locked) if markup else bar_plain(pct_of(f), f.get("status"), locked)
            rel = rel_reset(f.get("resets_at"), now)
        segs.append(f"{filled} {bar} {rel}")
    parts.append("  ".join(segs))
    bal = balance_str(account or {})
    fa = rel_ago(account.get("fetched_at") if account else None, now)
    parts.append(f"{bal} / {fa}")
    return "  ".join(parts)


def render_text(report, index, locked, now):
    accounts = (report or {}).get("accounts") or []
    if not accounts:
        return ("ocd: no data", worst_zone(None))
    idx = index % len(accounts)
    acct = accounts[idx]
    line = render_account_line(acct, now, locked, markup=True, with_name=True)
    return (line, worst_zone(acct, locked))


def render_tooltip(report, index, locked, now):
    accounts = (report or {}).get("accounts") or []
    if not accounts:
        return "ocd: waiting for first poll"
    gen = parse_dt((report or {}).get("generated_at"))
    head = "opencode limits"
    if gen:
        head += f"  \u00b7  updated {rel_ago(report.get('generated_at'), now)}"
    # Each window (5h/1w/1m) gets its own row; a blank row separates
    # accounts for a bigger gap. Waybar custom-module tooltips need \r
    # for line breaks (\n makes the tooltip vanish; see Waybar #3153).
    windows = (("usage_5h", "5h"), ("usage_week", "1w"), ("usage_month", "1m"))
    blocks = [head, ""]
    cur = index % len(accounts)
    for i, acct in enumerate(accounts):
        is_cur = i == cur
        marker = "\u25b8 " if is_cur else "  "
        row_locked = locked and is_cur
        nm = (acct or {}).get("account") or ""
        lock = " \U0001f512" if row_locked else ""
        bal = balance_str(acct or {})
        age = rel_ago((acct or {}).get("fetched_at"), now)
        blocks.append(f"{marker}{nm}{lock}   {bal} \u00b7 {age}")
        limits = (acct or {}).get("limits") or {}
        for key, wl in windows:
            f = frame_get(limits, key)
            if f is None:
                pct = 0.0
                status = None
                reset = "\u2014"
            else:
                pct = pct_of(f)
                status = f.get("status")
                reset = rel_reset(f.get("resets_at"), now)
            bar = bar_plain(pct, status, row_locked)
            blocks.append(f"  {wl}  [{int(round(pct)):>3}%] {bar}  {reset}")
        blocks.append("")
    return "\r".join(blocks)


def write_state(text, tooltip, klass):
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    payload = {
        "text": text,
        "tooltip": tooltip,
        "class": klass,
        "alt": "ocd",
    }
    tmp = STATE_FILE.with_suffix(".json.tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False))
    tmp.replace(STATE_FILE)


def fetch_ocd():
    """Run `ocd raw json`. Returns parsed dict or None on failure."""
    try:
        proc = subprocess.run(
            [OCD_BIN, "raw", "json"],
            capture_output=True, text=True, timeout=60,
        )
        if proc.returncode != 0:
            return None
        return json.loads(proc.stdout)
    except Exception:
        return None


class Daemon:
    def __init__(self):
        self.report = None
        self.index = 0
        self.last_advance = time.monotonic()
        self.locked_until = 0.0
        self.next_poll = 0.0
        self.last_text = None
        self.last_key = None
        self.sock = None
        self.running = True
        # The ocd poll runs in a worker thread so the main loop never
        # blocks on a slow network fetch. Rotation and clicks stay live.
        self._fetch_thread = None
        self._fresh_report = None
        self._fetch_lock = threading.Lock()

    def event_deadline(self):
        now = time.monotonic()
        deadlines = [now + TICK_SECS, self.next_poll]
        if not self.is_locked(now):
            deadlines.append(self.last_advance + ROTATE_SECS)
        else:
            deadlines.append(self.locked_until)
        return max(0.0, min(deadlines) - now)

    def is_locked(self, now=None):
        if now is None:
            now = time.monotonic()
        return self.locked_until and now < self.locked_until

    def advance(self):
        n = len((self.report or {}).get("accounts") or [])
        if n:
            self.index = (self.index + 1) % n
        self.last_advance = time.monotonic()

    def lock(self):
        self.locked_until = time.monotonic() + LOCK_SECS

    def _fetch_worker(self):
        rep = fetch_ocd()
        if rep:
            with self._fetch_lock:
                self._fresh_report = rep

    def maybe_tick(self):
        now = time.monotonic()
        # Kick off a poll on a worker thread; never block the main loop.
        if now >= self.next_poll and (self._fetch_thread is None or not self._fetch_thread.is_alive()):
            self.next_poll = now + POLL_SECS
            self._fetch_thread = threading.Thread(target=self._fetch_worker, daemon=True)
            self._fetch_thread.start()
        # Adopt a finished poll's result if one is ready.
        if self._fetch_thread is not None and not self._fetch_thread.is_alive():
            self._fetch_thread = None
            with self._fetch_lock:
                rep = self._fresh_report
                self._fresh_report = None
            if rep:
                self.report = rep
                n = len(rep.get("accounts") or [])
                if n and self.index >= n:
                    self.index %= n
        if not self.is_locked(now) and (now - self.last_advance) >= ROTATE_SECS:
            self.advance()
        if self.is_locked(now) and now >= self.locked_until:
            self.locked_until = 0.0

    def render_and_write(self):
        now = datetime.now(timezone.utc)
        text, zn = render_text(self.report, self.index, self.is_locked(), now)
        tooltip = render_tooltip(self.report, self.index, self.is_locked(), now)
        klass = ("ocd " + zn + (" ocd-locked" if self.is_locked() else "")).strip()
        # Write only when something changed. Displayed relative times
        # are minute-resolution, so the file updates a few times per minute.
        key = (text, tooltip, klass)
        if key != self.last_key:
            write_state(text, tooltip, klass)
            self.last_key = key

    def handle_cmd(self, cmd):
        cmd = (cmd or "").strip()
        if cmd == "advance":
            self.advance()
            self.locked_until = 0.0
        elif cmd == "lock":
            self.lock()
        elif cmd == "poll":
            self.next_poll = 0.0
        elif cmd == "quit":
            self.running = False

    def setup_socket(self):
        try:
            os.unlink(SOCK_FILE)
        except FileNotFoundError:
            pass
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        self.sock.bind(str(SOCK_FILE))
        self.sock.setblocking(False)
        os.chmod(SOCK_FILE, 0o600)

    def drain_socket(self):
        if self.sock is None:
            return
        while True:
            try:
                data, _ = self.sock.recvfrom(4096)
            except BlockingIOError:
                return
            self.handle_cmd(data.decode("utf-8", "replace"))

    def stop(self, *_):
        self.running = False
        try:
            if self.sock:
                self.sock.close()
            os.unlink(SOCK_FILE)
        except FileNotFoundError:
            pass
        sys.exit(0)


def main():
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    d = Daemon()
    d.setup_socket()
    signal.signal(signal.SIGTERM, d.stop)
    signal.signal(signal.SIGINT, d.stop)
    # force first poll
    d.next_poll = 0.0
    while d.running:
        d.maybe_tick()
        d.drain_socket()
        d.render_and_write()
        time.sleep(d.event_deadline())
    d.stop()


if __name__ == "__main__":
    main()