claude-usage/tray.py

#!/usr/bin/env python3
"""Claude Usage system tray indicator for Cinnamon."""

import gi
gi.require_version('Gtk', '3.0')
gi.require_version('AyatanaAppIndicator3', '0.1')

from gi.repository import Gtk, AyatanaAppIndicator3, GLib
import cairo
import json
import math
import os
import subprocess
import tempfile
import threading

SCRAPER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'scraper.ts')
BUN = os.path.expanduser('~/.bun/bin/bun')
ICON_DIR = os.path.join(tempfile.gettempdir(), 'claude-usage-icons')
LOCK_FILE = os.path.join(tempfile.gettempdir(), 'claude-usage-tray.pid')
ICON_SIZE = 48
POLL_SECONDS = 600  # 10 minutes
RETRY_SECONDS = 30  # retry quickly after failure


def acquire_lock():
    """Ensure only one instance runs. Kill ALL other tray.py processes."""
    import signal

    my_pid = os.getpid()

    # Kill any other running instances of this script (not just by PID file)
    try:
        result = subprocess.run(
            ['pgrep', '-f', 'claude-usage/tray.py'],
            capture_output=True, text=True,
        )
        for line in result.stdout.strip().split('\n'):
            if not line.strip():
                continue
            pid = int(line.strip())
            if pid != my_pid:
                try:
                    os.kill(pid, signal.SIGTERM)
                except (ProcessLookupError, PermissionError):
                    pass
    except Exception:
        pass

    # Also check PID file for good measure
    if os.path.exists(LOCK_FILE):
        try:
            old_pid = int(open(LOCK_FILE).read().strip())
            if old_pid != my_pid:
                os.kill(old_pid, signal.SIGTERM)
        except (ProcessLookupError, ValueError, PermissionError):
            pass

    import time
    time.sleep(0.5)

    with open(LOCK_FILE, 'w') as f:
        f.write(str(my_pid))


def release_lock():
    try:
        os.unlink(LOCK_FILE)
    except OSError:
        pass


class UsageIndicator:
    def __init__(self):
        self.data = None
        self.timer_id = None
        self.updating = False
        self.icon_counter = 0
        self.live_mode = False
        self.live_proc = None
        # Remember the real display so login browser opens on the user's screen
        self.real_display = os.environ.get('DISPLAY', ':0')

        os.makedirs(ICON_DIR, exist_ok=True)

        # Draw initial icon (grey, unknown state)
        initial_icon = self._draw_icon(None)
        initial_icon_name = os.path.splitext(os.path.basename(initial_icon))[0]

        self.indicator = AyatanaAppIndicator3.Indicator.new(
            'claude-usage',
            initial_icon_name,
            AyatanaAppIndicator3.IndicatorCategory.SYSTEM_SERVICES,
        )
        # Set the icon theme path so AppIndicator looks in our directory
        self.indicator.set_icon_theme_path(ICON_DIR)
        self.indicator.set_status(AyatanaAppIndicator3.IndicatorStatus.ACTIVE)

        self._build_menu()
        self.indicator.set_menu(self.menu)

        # Initial fetch
        self._schedule_fetch()
        # Start periodic timer
        self._start_timer()

    def _start_timer(self):
        if self.timer_id:
            GLib.source_remove(self.timer_id)
        self.timer_id = GLib.timeout_add_seconds(POLL_SECONDS, self._on_timer)

    def _on_timer(self):
        self._schedule_fetch()
        return True  # keep timer running

    def _schedule_fetch(self):
        if self.updating:
            return
        self.updating = True
        self._set_update_label('Updating...')
        t = threading.Thread(target=self._fetch_usage, daemon=True)
        t.start()

    def _fetch_usage(self):
        try:
            result = subprocess.run(
                [BUN, SCRAPER, '--json'],
                capture_output=True, text=True, timeout=120,
                env={**os.environ, 'NODE_NO_WARNINGS': '1'},
            )
            if result.returncode == 0 and result.stdout.strip():
                data = json.loads(result.stdout.strip())
                GLib.idle_add(self._on_data, data)
            else:
                self._log('Scraper failed: rc={} stdout={} stderr={}'.format(
                    result.returncode, result.stdout[:200], result.stderr[:200]))
                GLib.idle_add(self._on_data, {'loggedIn': False})
        except Exception as e:
            self._log('Fetch error: {}'.format(e))
            GLib.idle_add(self._on_data, {'loggedIn': False})

    @staticmethod
    def _log(msg):
        import datetime
        with open(os.path.join(tempfile.gettempdir(), 'claude-usage-tray.log'), 'a') as f:
            f.write('{} {}\n'.format(datetime.datetime.now().isoformat(), msg))

    def _on_data(self, data):
        self.data = data
        self.updating = False

        if data.get('loggedIn'):
            pct = data.get('sessionUsed', 0)
            icon_path = self._draw_icon(pct)
            self._log('OK: session={}% weekly={}%'.format(
                data.get('sessionUsed'), data.get('weeklyUsed')))
        else:
            icon_path = self._draw_icon_error()
            self._log('No data, will retry in {}s'.format(RETRY_SECONDS))
            # Schedule a fast retry instead of waiting 10 minutes
            GLib.timeout_add_seconds(RETRY_SECONDS, self._retry_after_failure)

        # set_icon_full wants icon NAME (no extension), looks in icon_theme_path
        icon_name = os.path.splitext(os.path.basename(icon_path))[0]
        self.indicator.set_icon_full(icon_name, 'Claude Usage {}%'.format(
            data.get('sessionUsed', '?') if data.get('loggedIn') else 'error'))
        self._update_menu()

    def _retry_after_failure(self):
        """Retry fetch after a short delay. Only retries if we still have no data."""
        if not self.data or not self.data.get('loggedIn'):
            self._schedule_fetch()
        return False  # don't repeat this timer

    def _next_icon_path(self):
        """Return a unique icon file path so AppIndicator detects changes."""
        self.icon_counter += 1
        return os.path.join(ICON_DIR, 'icon-{}.png'.format(self.icon_counter))

    def _draw_icon(self, pct):
        """Draw a circular progress indicator. pct=None for unknown state. Returns icon path."""
        icon_path = self._next_icon_path()
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, ICON_SIZE, ICON_SIZE)
        cr = cairo.Context(surface)
        cx, cy = ICON_SIZE / 2, ICON_SIZE / 2
        radius = ICON_SIZE / 2 - 3
        line_width = 5.0

        # Background ring
        cr.set_line_width(line_width)
        cr.set_source_rgba(0.3, 0.3, 0.3, 0.8)
        cr.arc(cx, cy, radius, 0, 2 * math.pi)
        cr.stroke()

        if pct is not None:
            # Progress arc (even at 0%, show green to distinguish from loading)
            if pct >= 90:
                cr.set_source_rgb(0.9, 0.2, 0.2)   # red
            elif pct >= 70:
                cr.set_source_rgb(0.9, 0.7, 0.1)    # yellow
            else:
                cr.set_source_rgb(0.2, 0.8, 0.3)    # green

            cr.set_line_width(line_width)
            cr.set_line_cap(cairo.LINE_CAP_ROUND)
            angle = (pct / 100.0) * 2 * math.pi
            if angle > 0.01:
                cr.arc(cx, cy, radius, -math.pi / 2, -math.pi / 2 + angle)
                cr.stroke()

            # Percentage text in center
            cr.set_source_rgb(0.9, 0.9, 0.9)
            cr.select_font_face('Sans', cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)
            cr.set_font_size(14)
            text = str(pct)
            extents = cr.text_extents(text)
            cr.move_to(cx - extents.width / 2, cy + extents.height / 2)
            cr.show_text(text)
        else:
            # Loading state: show "..." in center
            cr.set_source_rgba(0.6, 0.6, 0.6, 0.8)
            cr.select_font_face('Sans', cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)
            cr.set_font_size(14)
            text = '...'
            extents = cr.text_extents(text)
            cr.move_to(cx - extents.width / 2, cy + extents.height / 2)
            cr.show_text(text)

        surface.write_to_png(icon_path)
        return icon_path

    def _draw_icon_error(self):
        """Draw an X icon for logged-out / error state. Returns icon path."""
        icon_path = self._next_icon_path()
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, ICON_SIZE, ICON_SIZE)
        cr = cairo.Context(surface)
        cx, cy = ICON_SIZE / 2, ICON_SIZE / 2
        radius = ICON_SIZE / 2 - 3

        # Background ring
        cr.set_line_width(5.0)
        cr.set_source_rgba(0.3, 0.3, 0.3, 0.8)
        cr.arc(cx, cy, radius, 0, 2 * math.pi)
        cr.stroke()

        # Red X
        cr.set_source_rgb(0.9, 0.2, 0.2)
        cr.set_line_width(4.0)
        cr.set_line_cap(cairo.LINE_CAP_ROUND)
        offset = ICON_SIZE / 5
        cr.move_to(cx - offset, cy - offset)
        cr.line_to(cx + offset, cy + offset)
        cr.stroke()
        cr.move_to(cx + offset, cy - offset)
        cr.line_to(cx - offset, cy + offset)
        cr.stroke()

        surface.write_to_png(icon_path)
        return icon_path

    def _build_menu(self):
        self.menu = Gtk.Menu()

        self.item_session = Gtk.MenuItem(label='Session: --')
        self.item_session.set_sensitive(False)
        self.menu.append(self.item_session)

        self.item_session_reset = Gtk.MenuItem(label='')
        self.item_session_reset.set_sensitive(False)
        self.menu.append(self.item_session_reset)

        self.menu.append(Gtk.SeparatorMenuItem())

        self.item_weekly = Gtk.MenuItem(label='Weekly: --')
        self.item_weekly.set_sensitive(False)
        self.menu.append(self.item_weekly)

        self.item_weekly_reset = Gtk.MenuItem(label='')
        self.item_weekly_reset.set_sensitive(False)
        self.menu.append(self.item_weekly_reset)

        self.menu.append(Gtk.SeparatorMenuItem())

        self.item_update = Gtk.MenuItem(label='Update Now')
        self.item_update.connect('activate', self._on_update)
        self.menu.append(self.item_update)

        self.item_live = Gtk.MenuItem(label='Live Mode')
        self.item_live.connect('activate', self._on_toggle_live)
        self.menu.append(self.item_live)

        self.item_login = Gtk.MenuItem(label='Login')
        self.item_login.connect('activate', self._on_login)
        self.menu.append(self.item_login)

        self.menu.append(Gtk.SeparatorMenuItem())

        item_quit = Gtk.MenuItem(label='Quit')
        item_quit.connect('activate', lambda _: Gtk.main_quit())
        self.menu.append(item_quit)

        self.menu.show_all()

    def _update_menu(self):
        if not self.data or not self.data.get('loggedIn'):
            self.item_session.set_label('Session: not logged in')
            self.item_session_reset.set_label('')
            self.item_weekly.set_label('Weekly: --')
            self.item_weekly_reset.set_label('')
            self._set_update_label('Update Now')
            return

        d = self.data
        self.item_session.set_label(
            'Session: {}%'.format(d.get('sessionUsed', '?'))
        )
        self.item_session_reset.set_label(
            'Resets {}'.format(d.get('sessionResets', '?'))
        )
        self.item_weekly.set_label(
            'Weekly: {}%'.format(d.get('weeklyUsed', '?'))
        )
        self.item_weekly_reset.set_label(
            'Resets {}'.format(d.get('weeklyResets', '?'))
        )
        self._set_update_label('Update Now')

    def _set_update_label(self, text):
        self.item_update.set_label(text)

    def _on_update(self, _widget):
        self._start_timer()  # reset the 10-min timer
        self._schedule_fetch()

    def _on_toggle_live(self, widget):
        """Toggle live mode on/off."""
        if self.live_mode:
            self._stop_live()
        else:
            self._start_live()

    def _start_live(self):
        """Start live mode: keep Chromium open, read DOM every 15s."""
        if self.live_proc:
            return
        self.live_mode = True
        self._update_live_label()
        self._log('Live mode started')

        # Stop the normal polling timer
        if self.timer_id:
            GLib.source_remove(self.timer_id)
            self.timer_id = None

        def _run_live():
            try:
                self.live_proc = subprocess.Popen(
                    [BUN, SCRAPER, '--live'],
                    stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                    env={**os.environ, 'NODE_NO_WARNINGS': '1'},
                    bufsize=1,
                )
                for line in iter(self.live_proc.stdout.readline, b''):
                    line = line.decode('utf-8').strip()
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                        GLib.idle_add(self._on_data, data)
                    except json.JSONDecodeError:
                        pass
                # Process ended
                self.live_proc = None
                GLib.idle_add(self._on_live_stopped)
            except Exception as e:
                self._log('Live error: {}'.format(e))
                self.live_proc = None
                GLib.idle_add(self._on_live_stopped)

        t = threading.Thread(target=_run_live, daemon=True)
        t.start()

    def _stop_live(self):
        """Stop live mode, return to normal polling."""
        self.live_mode = False
        self._update_live_label()
        self._log('Live mode stopped')

        if self.live_proc:
            try:
                self.live_proc.terminate()
            except Exception:
                pass
            self.live_proc = None

        # Resume normal polling
        self._start_timer()
        self._schedule_fetch()

    def _on_live_stopped(self):
        """Called when live process exits unexpectedly."""
        if self.live_mode:
            self.live_mode = False
            self._update_live_label()
            self._log('Live process ended, resuming normal polling')
            self._start_timer()
            self._schedule_fetch()

    def _update_live_label(self):
        if self.live_mode:
            self.item_live.set_label('Live Mode [ON]')
        else:
            self.item_live.set_label('Live Mode')

    def _on_login(self, _widget):
        """Open browser for login on the user's real display, then auto-refresh."""
        def _do_login():
            env = {**os.environ, 'NODE_NO_WARNINGS': '1', 'DISPLAY': self.real_display}
            proc = subprocess.run(
                [BUN, SCRAPER, '--login'],
                capture_output=True, text=True, timeout=300,
                env=env,
            )
            # After login completes, fetch fresh usage data
            GLib.idle_add(self._schedule_fetch)

        t = threading.Thread(target=_do_login, daemon=True)
        t.start()


def main():
    import signal

    acquire_lock()

    indicator = UsageIndicator()

    def _shutdown(*args):
        # Kill live process if running
        if indicator.live_proc:
            try:
                indicator.live_proc.terminate()
            except Exception:
                pass
        # Hide the tray icon so Cinnamon removes it immediately
        indicator.indicator.set_status(AyatanaAppIndicator3.IndicatorStatus.PASSIVE)
        release_lock()
        Gtk.main_quit()

    signal.signal(signal.SIGTERM, _shutdown)
    signal.signal(signal.SIGINT, _shutdown)

    try:
        Gtk.main()
    finally:
        release_lock()


if __name__ == '__main__':
    main()