codex_status_print.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sqlite3
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from textwrap import shorten

from PIL import Image, ImageDraw, ImageFont


ROOT = Path(__file__).resolve().parent
CODEX_HOME = Path.home() / ".codex"
STATE_DB = CODEX_HOME / "state_5.sqlite"
GOALS_DB = CODEX_HOME / "goals_1.sqlite"
DEFAULT_OUT = ROOT / "codex_status.png"
PRINT_SCRIPT = ROOT / "s1_print.py"
DEFAULT_BLUETOOTH = os.environ.get("S1_BLUETOOTH_TARGET", "YOUR_PRINTER_NAME_OR_ADDRESS")
WIDTH = 384


def font(path: str, size: int):
    try:
        return ImageFont.truetype(path, size)
    except OSError:
        return ImageFont.load_default()


FONT_TITLE = font("C:/Windows/Fonts/consolab.ttf", 24)
FONT_BODY = font("C:/Windows/Fonts/consola.ttf", 15)
FONT_SMALL = font("C:/Windows/Fonts/consola.ttf", 13)
FONT_TINY = font("C:/Windows/Fonts/consola.ttf", 11)


def norm_path(value: str | None) -> str:
    if not value:
        return ""
    value = value.replace("\\\\?\\", "")
    return value.rstrip("\\/").lower()


def fmt_time(ms: int | None) -> str:
    if not ms:
        return "--"
    dt = datetime.fromtimestamp(ms / 1000, tz=timezone.utc).astimezone()
    return dt.strftime("%d %b %H:%M")


def fmt_sandbox(value: str | None) -> str:
    if not value:
        return "--"
    try:
        data = json.loads(value)
    except json.JSONDecodeError:
        return clip(value, 24)
    sandbox_type = data.get("type", "--")
    fs = data.get("file_system", {}).get("type")
    if fs:
        return f"{sandbox_type} / {fs}"
    network = data.get("network")
    if network:
        return f"{sandbox_type} / {network}"
    return sandbox_type


def clip(text: str, length: int) -> str:
    return shorten(text, width=length, placeholder="…")


def load_current_thread() -> dict | None:
    if not STATE_DB.exists():
        return None
    cwd = norm_path(str(ROOT))
    con = sqlite3.connect(STATE_DB)
    con.row_factory = sqlite3.Row
    try:
        rows = con.execute(
            """
            SELECT *
            FROM threads
            ORDER BY updated_at_ms DESC
            """
        ).fetchall()
    finally:
        con.close()
    for row in rows:
        if norm_path(row["cwd"]) == cwd:
            return dict(row)
    return dict(rows[0]) if rows else None


def load_active_goal(thread_id: str | None) -> dict | None:
    if not thread_id or not GOALS_DB.exists():
        return None
    con = sqlite3.connect(GOALS_DB)
    con.row_factory = sqlite3.Row
    try:
        rows = con.execute(
            """
            SELECT *
            FROM thread_goals
            WHERE thread_id = ?
            ORDER BY updated_at_ms DESC
            """,
            (thread_id,),
        ).fetchall()
    finally:
        con.close()
    active_states = {"active", "paused", "blocked", "usage_limited", "budget_limited"}
    for row in rows:
        if row["status"] in active_states:
            return dict(row)
    return dict(rows[0]) if rows else None


def draw_icon(d: ImageDraw.ImageDraw, x: int, y: int, s: int, kind: str) -> None:
    if kind == "chip":
        d.rectangle((x + 8, y + 8, x + s - 8, y + s - 8), outline=0, width=2)
        for off in (12, 20, 28):
            d.line((x + off, y + 4, x + off, y + 8), fill=0, width=1)
            d.line((x + off, y + s - 8, x + off, y + s - 4), fill=0, width=1)
            d.line((x + 4, y + off, x + 8, y + off), fill=0, width=1)
            d.line((x + s - 8, y + off, x + s - 4, y + off), fill=0, width=1)
        d.rectangle((x + 14, y + 14, x + s - 14, y + s - 14), outline=0, width=2)
    elif kind == "gauge":
        d.arc((x + 4, y + 4, x + s - 4, y + s - 4), 180, 0, fill=0, width=3)
        d.line((x + s // 2, y + s // 2, x + int(s * 0.75), y + int(s * 0.32)), fill=0, width=3)
        d.ellipse((x + s // 2 - 3, y + s // 2 - 3, x + s // 2 + 3, y + s // 2 + 3), fill=0)
    elif kind == "stack":
        for i in range(3):
            yy = y + 8 + i * 10
            d.rectangle((x + 8, yy, x + s - 8, yy + 6), outline=0, width=1)
    elif kind == "clock":
        d.ellipse((x + 5, y + 5, x + s - 5, y + s - 5), outline=0, width=3)
        cx = x + s // 2
        cy = y + s // 2
        d.line((cx, cy, cx, y + 16), fill=0, width=3)
        d.line((cx, cy, x + s - 18, cy + 2), fill=0, width=3)


def draw_key_value(d: ImageDraw.ImageDraw, x: int, y: int, key: str, value: str, width: int) -> None:
    d.text((x, y), key, fill=0, font=FONT_TINY)
    d.text((x + 82, y), clip(value, max(12, (width - 82) // 7)), fill=0, font=FONT_BODY)


def render_status(out: Path, thread: dict | None, goal: dict | None) -> None:
    lines = 9 if goal else 8
    height = 92 + lines * 24 + 26
    img = Image.new("L", (WIDTH, height), 255)
    d = ImageDraw.Draw(img)

    d.rectangle((0, 0, WIDTH - 1, height - 1), outline=0, width=2)
    draw_icon(d, 10, 10, 42, "chip")
    d.text((62, 11), "CODEX STATUS", fill=0, font=FONT_TITLE)
    d.text((62, 38), datetime.now().astimezone().strftime("%a %d %b %H:%M"), fill=0, font=FONT_SMALL)
    d.line((8, 58, WIDTH - 9, 58), fill=0, width=2)

    y = 70
    d.text((12, y), "THREAD", fill=0, font=FONT_TINY)
    d.text((86, y), clip(thread.get("title", "unknown") if thread else "unknown", 31), fill=0, font=FONT_BODY)
    y += 24
    draw_key_value(d, 12, y, "MODEL", thread.get("model", "--") if thread else "--", WIDTH - 24)
    y += 22
    draw_key_value(d, 12, y, "EFFORT", thread.get("reasoning_effort", "--") if thread else "--", WIDTH - 24)
    y += 22
    draw_key_value(d, 12, y, "TOKENS", f"{thread.get('tokens_used', 0):,}" if thread else "0", WIDTH - 24)
    y += 22
    draw_key_value(d, 12, y, "SANDBOX", fmt_sandbox(thread.get("sandbox_policy")) if thread else "--", WIDTH - 24)
    y += 22
    draw_key_value(d, 12, y, "APPROVAL", thread.get("approval_mode", "--") if thread else "--", WIDTH - 24)
    y += 22
    draw_key_value(d, 12, y, "UPDATED", fmt_time(thread.get("updated_at_ms")) if thread else "--", WIDTH - 24)
    y += 24

    d.line((8, y - 4, WIDTH - 9, y - 4), fill=0, width=2)
    draw_icon(d, 10, y + 2, 28, "gauge")
    if goal:
        budget = goal.get("token_budget")
        used = goal.get("tokens_used", 0)
        status = goal.get("status", "unknown")
        objective = goal.get("objective", "")
        d.text((50, y + 2), clip(objective, 36), fill=0, font=FONT_SMALL)
        y += 20
        draw_key_value(d, 12, y + 2, "GOAL", status, WIDTH - 24)
        y += 22
        budget_text = f"{used:,}"
        if budget:
            budget_text = f"{used:,}/{budget:,}"
        draw_key_value(d, 12, y + 2, "USAGE", budget_text, WIDTH - 24)
        y += 22
        draw_key_value(d, 12, y + 2, "GOAL AT", fmt_time(goal.get("updated_at_ms")), WIDTH - 24)
        y += 22
    else:
        d.text((50, y + 2), "NO ACTIVE GOAL", fill=0, font=FONT_SMALL)
        y += 20
        draw_key_value(d, 12, y + 2, "GOAL", "idle", WIDTH - 24)
        y += 22
        draw_key_value(d, 12, y + 2, "USAGE", "session only", WIDTH - 24)
        y += 22

    d.line((8, height - 26, WIDTH - 9, height - 26), fill=0, width=2)
    d.text((10, height - 20), clip(str(ROOT), 50), fill=0, font=FONT_TINY)
    img.save(out)


def print_image(path: Path, darkness: int) -> int:
    command = [
        str(ROOT / ".venv" / "Scripts" / "python.exe"),
        str(PRINT_SCRIPT),
        "image",
        str(path),
        "--darkness",
        str(darkness),
    ]
    return subprocess.call(command, cwd=ROOT)


def main() -> int:
    parser = argparse.ArgumentParser(description="Print Codex thread status and usage.")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--darkness", type=int, choices=range(1, 6), default=3)
    parser.add_argument("--bluetooth", default=DEFAULT_BLUETOOTH)
    parser.add_argument("--no-print", action="store_true")
    args = parser.parse_args()

    thread = load_current_thread()
    goal = load_active_goal(None if thread is None else thread.get("id"))
    render_status(args.out, thread, goal)
    print(f"Wrote {args.out}")
    if args.no_print:
        return 0
    return print_image(args.out, args.darkness)


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