token_ghost.py

#!/usr/bin/env python3
"""Token Ghost collectors + SwiftBar renderer.

Sources:
- Claude Code: interactive `/usage` capture. Claude Code keeps no structured
  usage-limit data on disk, so a short TUI session is the only local source.
  Results are cached and reused until CLAUDE_TTL_SECONDS passes.
- Codex: structured `rate_limits` events in local session logs
  (~/.codex/sessions/**/*.jsonl). These update only when Codex actually runs,
  so the UI shows the data age and marks a bucket as "reset (est.)" once its
  reset time has passed.

No credentials are printed. Cache lives at ~/.cache/token-ghost/cache.json.
"""
from __future__ import annotations

import json
import os
import pty
import re
import select
import shutil
import signal
import struct
import sys
import termios
import time
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
import fcntl

CACHE_PATH = Path.home() / ".cache" / "token-ghost" / "cache.json"

# Capturing Claude /usage spawns an interactive session (~15s), too heavy for
# every 5-minute SwiftBar tick. Cached values are reused until this TTL passes;
# the Refresh menu action bypasses it with --force.
try:
    CLAUDE_TTL_SECONDS = float(os.environ.get("TOKEN_GHOST_CLAUDE_TTL_MIN", "15")) * 60
except ValueError:
    CLAUDE_TTL_SECONDS = 15 * 60

# SwiftBar launches plugins with a minimal PATH, unlike the user's Terminal.
# Add the common macOS Node/Homebrew locations so `claude` and `codex` can be found.
COMMON_BIN_DIRS = [
    "/opt/homebrew/bin",
    "/usr/local/bin",
    "/usr/bin",
    "/bin",
    "/usr/sbin",
    "/sbin",
    str(Path.home() / ".npm-global" / "bin"),
    str(Path.home() / ".local" / "bin"),
]
# nvm installs live under versioned directories; include them so nvm-managed
# `claude`/`codex` binaries are found too (newest version first).
COMMON_BIN_DIRS.extend(
    str(p) for p in sorted(Path.home().glob(".nvm/versions/node/*/bin"), reverse=True)
)
_existing_path = os.environ.get("PATH", "")
os.environ["PATH"] = ":".join([p for p in COMMON_BIN_DIRS + _existing_path.split(":") if p])

ANSI_RE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]|\x1b\][^\x07]*(?:\x07|\x1b\\)|\x1b[()][A-Za-z0-9]")


@dataclass
class ClaudeSession:
    used_percent: int | None = None
    reset_text: str | None = None
    timezone: str | None = None


@dataclass
class ClaudeExtra:
    used_percent: int | None = None
    spent_usd: float | None = None
    limit_usd: float | None = None
    reset_text: str | None = None
    timezone: str | None = None


@dataclass
class UsageBucket:
    used_percent: int | None = None
    reset_text: str | None = None
    timezone: str | None = None


@dataclass
class CodexFiveHour:
    left_percent: int | None = None
    used_percent: int | None = None
    reset_text: str | None = None


@dataclass
class CodexWeek:
    left_percent: int | None = None
    used_percent: int | None = None
    reset_text: str | None = None


def strip_ansi(text: str) -> str:
    return ANSI_RE.sub("", text).replace("\r", "\n")


def compact(text: str) -> str:
    # TUI captures sometimes lose or distort spaces. Normalize by removing whitespace.
    return re.sub(r"\s+", "", strip_ansi(text))


def normalize_reset_text(value: str | None) -> str | None:
    if value is None:
        return None
    value = value.strip()
    value = re.sub(r"([A-Za-z])([0-9])", r"\1 \2", value)
    value = re.sub(r"([0-9])([A-Za-z])", r"\1 \2", value)
    return value


def parse_claude_usage(text: str) -> dict:
    c = compact(text)
    session = ClaudeSession()
    extra = ClaudeExtra()
    week_all = UsageBucket()

    m = re.search(r"Currentsession.*?(\d{1,3})%used.*?Resets([^()]+)\(([^()]+)\)", c, re.S)
    if m:
        session.used_percent = int(m.group(1))
        session.reset_text = normalize_reset_text(m.group(2))
        session.timezone = m.group(3)

    m = re.search(r"Currentweek\(allmodels\).*?(\d{1,3})%used.*?Resets([^()]+)\(([^()]+)\)", c, re.S)
    if m:
        week_all.used_percent = int(m.group(1))
        week_all.reset_text = normalize_reset_text(m.group(2))
        week_all.timezone = m.group(3)

    m = re.search(
        r"Extrausage.*?(\d{1,3})%used.*?\$([0-9.,]+)/\$([0-9.,]+)spent·?Resets([^()]+)\(([^()]+)\)",
        c,
        re.S,
    )
    if m:
        extra.used_percent = int(m.group(1))
        extra.spent_usd = float(m.group(2).replace(",", ""))
        extra.limit_usd = float(m.group(3).replace(",", ""))
        extra.reset_text = normalize_reset_text(m.group(4))
        extra.timezone = m.group(5)

    return {"session": asdict(session), "week_all": asdict(week_all), "extra": asdict(extra)}


def _drain(fd: int, seconds: float) -> bytes:
    out = b""
    end = time.time() + seconds
    while time.time() < end:
        readable, _, _ = select.select([fd], [], [], 0.1)
        if fd not in readable:
            continue
        try:
            chunk = os.read(fd, 8192)
        except OSError:
            break
        if not chunk:
            break
        out += chunk
    return out


def run_tui(command: list[str], slash_command: str, startup_wait: float = 3, after_wait: float = 8, pre_enter: bool = False) -> str:
    pid, fd = pty.fork()
    if pid == 0:
        os.environ.setdefault("TERM", "xterm-256color")
        os.environ.setdefault("COLUMNS", "120")
        os.environ.setdefault("LINES", "40")
        try:
            fcntl.ioctl(1, termios.TIOCSWINSZ, struct.pack("HHHH", 40, 120, 0, 0))
        except Exception:
            pass
        os.execvp(command[0], command)

    output = b""
    try:
        fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", 40, 120, 0, 0))
    except Exception:
        pass

    try:
        output += _drain(fd, startup_wait)
        if pre_enter:
            # Claude first-run workspace trust dialog defaults to Yes.
            os.write(fd, b"\r")
            output += _drain(fd, 3)
        os.write(fd, b"\x15")  # Ctrl-U: clear starter prompt
        output += _drain(fd, 0.3)
        os.write(fd, slash_command.encode("utf-8") + b"\r")
        output += _drain(fd, after_wait)
        try:
            os.write(fd, b"\x03")
            output += _drain(fd, 0.5)
        except OSError:
            pass
    finally:
        try:
            os.kill(pid, signal.SIGTERM)
        except Exception:
            pass
    return output.decode("utf-8", "replace")


def collect_claude() -> dict:
    if not shutil.which("claude"):
        return {"error": "claude CLI not found"}
    text = run_tui(["claude"], "/usage", startup_wait=3, after_wait=10, pre_enter=True)
    parsed = parse_claude_usage(text)
    ok = parsed["session"]["used_percent"] is not None or parsed["extra"]["used_percent"] is not None
    if not ok:
        parsed["error"] = "Claude /usage parse failed"
    return parsed


def _epoch_is_past(value: int | float | None) -> bool:
    if value is None:
        return False
    try:
        return float(value) < time.time() - 60
    except Exception:
        return False


def _format_reset_epoch(value: int | float | None, weekly: bool = False) -> str | None:
    if value is None:
        return None
    try:
        fmt = "%b %-d %-I%p" if weekly else "%-I:%M%p"
        return time.strftime(fmt, time.localtime(float(value))).lower()
    except Exception:
        return None


def _project_reset_epoch(resets_at: int | float | None, window_minutes: int | float | None) -> float | None:
    """When a reset time has passed, roll it forward by whole windows to the
    next upcoming reset. Codex reports resets_at for the *current* window and a
    fixed window length, so the next reset is resets_at + k*window. This is an
    estimate: the real next reset is set when Codex next talks to the backend."""
    if resets_at is None:
        return None
    try:
        value = float(resets_at)
    except (TypeError, ValueError):
        return None
    now = time.time()
    if value >= now or not window_minutes:
        return value
    try:
        window = float(window_minutes) * 60
    except (TypeError, ValueError):
        return value
    if window <= 0:
        return value
    steps = (now - value) // window + 1
    return value + steps * window


def _parse_iso_timestamp(value: str | None) -> float | None:
    if not value:
        return None
    try:
        parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
        if parsed.tzinfo is None:
            parsed = parsed.replace(tzinfo=timezone.utc)
        return parsed.timestamp()
    except Exception:
        return None


def age_text(epoch: float | None) -> str | None:
    if not epoch:
        return None
    seconds = max(0.0, time.time() - epoch)
    if seconds < 90:
        return "just now"
    minutes = seconds / 60
    if minutes < 60:
        return f"{int(minutes)}m ago"
    hours = minutes / 60
    if hours < 48:
        return f"{int(hours)}h ago"
    return f"{int(hours / 24)}d ago"


def parse_codex_rate_limits(rate_limits: dict | None) -> dict:
    five = CodexFiveHour()
    week_all = CodexWeek()
    if not isinstance(rate_limits, dict):
        return {"five_hour": asdict(five), "week_all": asdict(week_all)}

    primary = rate_limits.get("primary") or {}
    secondary = rate_limits.get("secondary") or {}

    primary_reset_expired = _epoch_is_past(primary.get("resets_at"))
    secondary_reset_expired = _epoch_is_past(secondary.get("resets_at"))

    # When a window has passed, show the projected next reset instead of a stale
    # past time. When it is still valid, this returns the reported reset as-is.
    primary_reset = _project_reset_epoch(primary.get("resets_at"), primary.get("window_minutes"))
    secondary_reset = _project_reset_epoch(secondary.get("resets_at"), secondary.get("window_minutes"))

    if primary.get("used_percent") is not None:
        five.used_percent = int(round(float(primary["used_percent"])))
        five.left_percent = max(0, 100 - five.used_percent)
        five.reset_text = _format_reset_epoch(primary_reset, weekly=False)

    if secondary.get("used_percent") is not None:
        week_all.used_percent = int(round(float(secondary["used_percent"])))
        week_all.left_percent = max(0, 100 - week_all.used_percent)
        week_all.reset_text = _format_reset_epoch(secondary_reset, weekly=True)

    result = {"five_hour": asdict(five), "week_all": asdict(week_all)}
    result["five_hour"]["reset_expired"] = primary_reset_expired
    result["week_all"]["reset_expired"] = secondary_reset_expired
    return result


def _tail_lines(path: Path, max_bytes: int = 65536) -> list[str]:
    """Read only the end of a session log. rate_limits events land near the
    tail, and full session files can be tens of megabytes."""
    try:
        with path.open("rb") as fh:
            fh.seek(0, os.SEEK_END)
            size = fh.tell()
            fh.seek(max(0, size - max_bytes))
            data = fh.read()
    except Exception:
        return []
    return data.decode("utf-8", "replace").splitlines()


def collect_codex() -> dict:
    """Read Codex usage from local session logs — the single source.

    Codex records structured rate_limits (5h primary + weekly secondary,
    used_percent / resets_at) whenever it talks to the backend. TUI automation
    (/status scraping) was removed: update prompts, MCP startup spinners and
    locale differences made it unreliable. The trade-off is that values only
    move when Codex actually runs, which the UI communicates via data age and
    "reset (est.)" once a bucket's reset time passes.
    """
    def mtime(path: Path) -> float:
        try:
            return path.stat().st_mtime
        except OSError:
            return 0.0

    session_paths = sorted(
        Path.home().glob(".codex/sessions/**/*.jsonl"),
        key=mtime,
        reverse=True,
    )
    cli_missing = shutil.which("codex") is None
    newest: dict | None = None
    newest_epoch = 0.0

    # Selection is purely "newest event by timestamp". Bucket expiry is a
    # rendering concern (per bucket), not a reason to prefer another event.
    for path in session_paths[:100]:
        # A file cannot contain events newer than its mtime, so once the event
        # in hand is newer than the remaining files, we are done (60s margin).
        if newest and mtime(path) < newest_epoch - 60:
            break
        for line in reversed(_tail_lines(path)):
            try:
                obj = json.loads(line)
            except Exception:
                continue
            payload = obj.get("payload", {}) if isinstance(obj, dict) else {}
            info = payload.get("info", {}) if isinstance(payload, dict) else {}
            rate_limits = info.get("rate_limits") or payload.get("rate_limits")
            if not rate_limits:
                continue
            parsed = parse_codex_rate_limits(rate_limits)
            if parsed["five_hour"]["used_percent"] is None:
                continue
            timestamp = obj.get("timestamp") or ""
            epoch = _parse_iso_timestamp(timestamp) or 0.0
            if newest is None or epoch > newest_epoch:
                parsed["source"] = "local Codex session log"
                parsed["source_timestamp"] = timestamp
                parsed["source_epoch"] = epoch or None
                newest = parsed
                newest_epoch = epoch
            # The first rate_limits event found scanning backwards is the
            # newest in this file; older lines cannot beat it.
            break

    if newest:
        newest["cli_missing"] = cli_missing
        return newest

    parsed = parse_codex_rate_limits(None)
    parsed["cli_missing"] = cli_missing
    parsed["error"] = (
        "Codex CLI not found" if cli_missing else "Codex usage not found in local sessions yet"
    )
    return parsed


def load_cache() -> dict:
    try:
        return json.loads(CACHE_PATH.read_text())
    except Exception:
        return {}


def save_cache(data: dict) -> None:
    CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
    tmp = CACHE_PATH.with_suffix(".tmp")
    tmp.write_text(json.dumps(data, indent=2, ensure_ascii=False))
    tmp.replace(CACHE_PATH)


def _claude_has_data(claude: dict) -> bool:
    return (
        claude.get("session", {}).get("used_percent") is not None
        or claude.get("extra", {}).get("used_percent") is not None
    )


def collect_claude_cached(cache: dict, force: bool = False) -> dict:
    """Return Claude usage, spawning the /usage capture only when the cached
    value is missing or older than CLAUDE_TTL_SECONDS (or on --force)."""
    now = time.time()
    cached = cache.get("claude") or {}
    cached_age = now - float(cached.get("collected_at") or 0)
    if not force and _claude_has_data(cached) and cached_age < CLAUDE_TTL_SECONDS:
        return cached

    claude = collect_claude()
    if _claude_has_data(claude):
        claude["collected_at"] = now
        return claude
    if _claude_has_data(cached):
        # Keep the last good value (and its collected_at, so the next tick
        # retries) instead of blanking the menu on one failed capture.
        return cached | {"error": claude.get("error") or "Claude usage refresh failed"}
    return claude


def collect_all(force: bool = False) -> dict:
    cache = load_cache()
    result = {
        "updated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
        "claude": collect_claude_cached(cache, force=force),
        "codex": collect_codex(),
    }
    save_cache(result)
    return result


def severity_claude(used: int | None) -> str:
    if used is None:
        return "gray"
    if used >= 90:
        return "red"
    if used >= 70:
        return "orange"
    return "green"


def severity_codex_left(left: int | None) -> str:
    if left is None:
        return "gray"
    if left <= 19:
        return "red"
    if left <= 49:
        return "orange"
    return "green"


def pct(v: int | None, suffix: str = "%") -> str:
    return "?" if v is None else f"{v}{suffix}"


def title_pct(used: int | None) -> str:
    return "--" if used is None else f"{used}%"


def top_title(used: int | None) -> str:
    return f"👻{title_pct(used)}"


def progress_bar(used: int | None, width: int = 8) -> str:
    """High-contrast emoji progress bar for SwiftBar's light menu.

    Text block characters (█/░) rendered too similarly in SwiftBar, especially at
    low percentages. Emoji squares stay visibly colored without relying on font
    weight or row color.
    """
    if used is None:
        return "⬜️" * width
    used = max(0, min(100, int(used)))
    filled = 0 if used == 0 else max(1, min(width, round(used * width / 100)))
    if used >= 90:
        fill = "🟥"
    elif used >= 70:
        fill = "🟧"
    elif used >= 45:
        fill = "🟨"
    else:
        fill = "🟩"
    return fill * filled + "⬜️" * (width - filled)


def compact_reset_text(reset_text: str | None) -> str | None:
    """Render reset times in a short, consistent menu-bar format.

    Examples:
    - "3:40 pm" -> "3:40pm"
    - "Jul 15 at 6 pm" -> "Jul15 6pm"
    - "jul 14 10am" -> "Jul14 10am"
    """
    if not reset_text:
        return None
    text = reset_text.strip()
    text = re.sub(r"\s+", " ", text)
    text = re.sub(r"\s*([ap])\.?(m)\.?$", r"\1\2", text, flags=re.I)
    text = re.sub(r"(\d)\s+([ap]m)\b", r"\1\2", text, flags=re.I)
    text = re.sub(r"\s+at\s+", " ", text, flags=re.I)

    month_match = re.match(r"^([A-Za-z]{3,9})\s+(\d{1,2})(?:\s+(.+))?$", text)
    if month_match:
        month = month_match.group(1)[:3].title()
        day = month_match.group(2)
        tail = month_match.group(3) or ""
        tail = re.sub(r"(\d)\s+([ap]m)\b", r"\1\2", tail, flags=re.I).lower()
        return f"{month}{day}" + (f" {tail}" if tail else "")

    return text.lower()


def reset_suffix(reset_text: str | None) -> str:
    compacted = compact_reset_text(reset_text)
    return f" (~{compacted})" if compacted else ""


def money_suffix(spent: float | None, limit: float | None) -> str:
    if spent is None or limit is None:
        return ""
    return f" (${spent:.2f} / ${limit:.2f})"


def style(color: str = "#111111", size: int = 12, font: str | None = None, active: bool = True) -> str:
    parts = [f"color={color}", f"size={size}"]
    if font:
        parts.append(f"font={font}")
    # SwiftBar dims non-action rows in some macOS appearances. A no-op bash action
    # keeps informational rows rendered as normal dark text without opening Terminal.
    if active:
        parts.extend(["bash='/usr/bin/true'", "terminal=false"])
    return " | " + " ".join(parts)


def usage_line(label: str, used: int | None, suffix: str = "", unknown_text: str | None = None) -> str:
    if used is None and unknown_text:
        return f"{label}: {unknown_text}" + style(size=12)
    return f"{label}: {progress_bar(used)}  {pct(used)}{suffix}" + style(size=12)


def codex_usage_line(label: str, bucket: dict) -> str:
    used = bucket.get("used_percent")
    if used is None:
        return f"{label}: Not connected" + style(size=12)
    if bucket.get("reset_expired"):
        # The window reset after the last Codex run, so real usage is unknown
        # but presumably back to 0, and reset_text is the projected next reset.
        # Rendered in the same "(~time)" form as every other bucket; the 0% plus
        # the "From session log" age line already signal it is an estimate.
        return usage_line(label, 0, reset_suffix(bucket.get("reset_text")))
    return usage_line(label, used, reset_suffix(bucket.get("reset_text")))


def data_age_line(text: str) -> str:
    return text + style(color="#666666", size=11, active=False)


def setup_hint(tool: str, error: str | None) -> list[str]:
    """User-facing setup/help rows for missing or unreadable CLI state."""
    if not error:
        return []
    if tool == "claude":
        return [
            "Claude setup needed" + style(color="orange", size=12),
            "To show Claude usage:" + style(size=12),
            "1. Install Claude Code CLI" + style(size=12),
            "2. Open it once and log in (/login)" + style(size=12),
            "Run Claude setup in Terminal | bash='/bin/sh' param1='-lc' param2='command -v claude >/dev/null || npm install -g @anthropic-ai/claude-code; echo \"If Claude Code asks, log in with /login, then quit.\"; claude' terminal=true",
            "Copy install command | bash='/bin/sh' param1='-lc' param2='printf %s \"npm install -g @anthropic-ai/claude-code\" | pbcopy' terminal=false",
        ]
    return [
        "Codex setup needed" + style(color="orange", size=12),
        "To show Codex usage:" + style(size=12),
        "1. Install Codex CLI and log in" + style(size=12),
        "2. Run any Codex task once" + style(size=12),
        "(usage is read from local Codex session logs)" + style(size=12),
    ]


def render_swiftbar(data: dict) -> str:
    claude = data.get("claude", {})
    session = claude.get("session", {})
    claude_week = claude.get("week_all", {})
    extra = claude.get("extra", {})
    codex = data.get("codex", {})
    five = codex.get("five_hour", {})
    codex_week = codex.get("week_all", {})

    current = session.get("used_percent")
    title = top_title(current)

    lines = [
        title,
        "---",
        "🧡 Claude" + style(color="#111111", size=12),
        usage_line("Current", current, reset_suffix(session.get("reset_text")), "Not connected"),
        usage_line("Weekly", claude_week.get("used_percent"), reset_suffix(claude_week.get("reset_text")), "Not connected"),
    ]
    if extra.get("used_percent") is not None:
        lines.append(usage_line("Extra", extra.get("used_percent"), money_suffix(extra.get("spent_usd"), extra.get("limit_usd"))))
    claude_age = age_text(claude.get("collected_at"))
    if claude_age:
        lines.append(data_age_line(f"Checked {claude_age}"))
    lines.extend(setup_hint("claude", claude.get("error")))

    lines.append("💙 Codex" + style(color="#111111", size=12))
    if codex.get("error") and five.get("used_percent") is None:
        lines.extend(setup_hint("codex", codex.get("error")))
    else:
        lines.append(codex_usage_line("5 Hour", five))
        lines.append(codex_usage_line("Weekly", codex_week))
        codex_age = age_text(codex.get("source_epoch"))
        if codex_age:
            lines.append(data_age_line(f"From session log, {codex_age}"))
        if codex.get("cli_missing"):
            lines.append("Codex CLI not found — showing last session log" + style(color="orange", size=11, active=False))
        lines.extend(setup_hint("codex", codex.get("error")))

    script_path = str(Path.home() / ".token-ghost" / "token_ghost.py")
    lines += [
        "---",
        f"Last checked: {data.get('updated_at', 'unknown')}" + style(color="#666666", size=12, active=False),
        f"Refresh | bash='{script_path}' param1='--collect' param2='--force' terminal=false refresh=true",
        "Setup guide | bash='/bin/sh' param1='-c' param2='open \"$HOME/.token-ghost/README.md\" 2>/dev/null || open https://github.com/zoeymakes/token-ghost#install' terminal=false",
        "Open Claude | bash='/bin/sh' param1='-c' param2='open -a Claude 2>/dev/null || open https://claude.ai' terminal=false",
        "Open Codex | bash='/bin/sh' param1='-c' param2='open -a Codex 2>/dev/null || open -a ChatGPT 2>/dev/null || open https://chatgpt.com/codex' terminal=false",
    ]
    return "\n".join(lines)


def main(argv: list[str]) -> int:
    force = "--force" in argv
    if "--collect" in argv:
        print(json.dumps(collect_all(force=force), indent=2, ensure_ascii=False))
        return 0
    if "--render-cache" in argv:
        print(render_swiftbar(load_cache()))
        return 0
    print(render_swiftbar(collect_all(force=force)))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))