widget.py

"""
widget.py - Small always-on-top overlay widget for the Claude Token Counter.

Sits at the bottom-right corner of the screen, above the taskbar.
"""

import tkinter as tk
import tkinter.font as tkfont
import tkinter.messagebox as messagebox


# ---------------------------------------------------------------------------
# Color / style constants
# ---------------------------------------------------------------------------
BG_COLOR       = "#1a1a2e"
FG_COLOR       = "#e0e0e0"
LABEL_COLOR    = "#888888"
ACCENT_COLOR   = "#4fc3f7"
HIGHLIGHT_BG   = "#2d2d44"
DOT_ACTIVE     = "#4caf50"   # green  – proxy running
DOT_IDLE       = "#757575"   # gray   – idle
BORDER_COLOR   = "#2d2d44"

WIDGET_W       = 220
WIDGET_H       = 95
MARGIN_RIGHT   = 10
MARGIN_BOTTOM  = 50          # above taskbar


def _fmt(n):
    """Format a token count compactly: 1,234 or 2.3M."""
    try:
        n = int(n)
    except (TypeError, ValueError):
        return "0"
    if n >= 1_000_000:
        m = n / 1_000_000
        return f"{m:.1f}M" if m < 10 else f"{m:.0f}M"
    return f"{n:,}"


class TokenOverlay:
    """
    Small borderless overlay window that floats above the taskbar at the
    bottom-right of the screen.

    Parameters
    ----------
    root : tk.Tk
        The root Tk instance (may be iconified / withdrawn).
    storage : TokenStorage
        The shared storage backend.
    on_click_callback : callable, optional
        Called (with no arguments) when the widget is left-clicked to open
        the detailed view.
    refresh_callback : callable, optional
        Called every 2 seconds so the application can push fresh data.
    """

    def __init__(self, root, storage, on_click_callback=None, refresh_callback=None):
        self._root            = root
        self._storage         = storage
        self._on_click        = on_click_callback
        self._refresh_cb      = refresh_callback
        self._on_reset        = None

        # Drag state
        self._drag_x = 0
        self._drag_y = 0

        # Build the Toplevel window
        self._win = tk.Toplevel(root)
        self._win.overrideredirect(True)     # borderless
        self._win.attributes("-topmost", True)
        self._win.attributes("-alpha", 0.92)
        self._win.configure(bg=BG_COLOR)
        self._win.resizable(False, False)

        # Position at bottom-right
        self._position_window()

        # Fonts
        mono_font   = tkfont.Font(family="Consolas", size=10, weight="bold")
        label_font  = tkfont.Font(family="Consolas", size=7)
        session_font= tkfont.Font(family="Segoe UI",  size=7)

        # ----------------------------------------------------------------
        # Outer frame with a subtle border
        # ----------------------------------------------------------------
        outer = tk.Frame(
            self._win,
            bg=BORDER_COLOR,
            padx=1, pady=1
        )
        outer.pack(fill=tk.BOTH, expand=True)

        inner = tk.Frame(outer, bg=BG_COLOR, padx=8, pady=6)
        inner.pack(fill=tk.BOTH, expand=True)

        # ----------------------------------------------------------------
        # Row 0 – header row: "Claude Tokens" label + status dot
        # ----------------------------------------------------------------
        header_frame = tk.Frame(inner, bg=BG_COLOR)
        header_frame.pack(fill=tk.X)

        self._label_title = tk.Label(
            header_frame,
            text="G4 Claw Counter",
            font=label_font,
            bg=BG_COLOR,
            fg=LABEL_COLOR,
            anchor="w"
        )
        self._label_title.pack(side=tk.LEFT)

        self._dot = tk.Label(
            header_frame,
            text="●",
            font=tkfont.Font(family="Segoe UI", size=8),
            bg=BG_COLOR,
            fg=DOT_IDLE
        )
        self._dot.pack(side=tk.RIGHT)

        # ----------------------------------------------------------------
        # Row 1 – token counts
        # ----------------------------------------------------------------
        self._label_counts = tk.Label(
            inner,
            text="In: 0 | Out: 0",
            font=mono_font,
            bg=BG_COLOR,
            fg=FG_COLOR,
            anchor="w"
        )
        self._label_counts.pack(fill=tk.X, pady=(2, 0))

        # ----------------------------------------------------------------
        # Row 2 – today's cost
        # ----------------------------------------------------------------
        self._label_cost = tk.Label(
            inner,
            text="Today: $0.00",
            font=session_font,
            bg=BG_COLOR,
            fg=ACCENT_COLOR,
            anchor="w"
        )
        self._label_cost.pack(fill=tk.X, pady=(1, 0))

        # ----------------------------------------------------------------
        # Row 3 – session name
        # ----------------------------------------------------------------
        self._label_session = tk.Label(
            inner,
            text="0 models | 0 reqs today",
            font=session_font,
            bg=BG_COLOR,
            fg=LABEL_COLOR,
            anchor="w"
        )
        self._label_session.pack(fill=tk.X, pady=(1, 0))

        # ----------------------------------------------------------------
        # Bind events
        # ----------------------------------------------------------------
        for widget in (self._win, outer, inner,
                       header_frame, self._label_title, self._dot,
                       self._label_counts, self._label_cost, self._label_session):
            widget.bind("<Button-1>",        self._on_left_click)
            widget.bind("<Button-3>",        self._show_context_menu)
            widget.bind("<ButtonPress-1>",   self._drag_start)
            widget.bind("<B1-Motion>",       self._drag_motion)

        # ----------------------------------------------------------------
        # Context menu
        # ----------------------------------------------------------------
        self._context_menu = tk.Menu(
            self._win,
            tearoff=0,
            bg=HIGHLIGHT_BG,
            fg=FG_COLOR,
            activebackground=ACCENT_COLOR,
            activeforeground=BG_COLOR,
            bd=0,
            relief=tk.FLAT
        )
        self._context_menu.add_command(label="Detailed View",    command=self._open_detailed)
        self._context_menu.add_command(label="Reset All Data",  command=self._reset_session)
        self._context_menu.add_separator()
        self._context_menu.add_command(label="Quit",            command=self._quit_app)

        # ----------------------------------------------------------------
        # Start the auto-refresh loop
        # ----------------------------------------------------------------
        self._running = True
        self._schedule_refresh()

    # ------------------------------------------------------------------
    # Positioning helpers
    # ------------------------------------------------------------------

    def _position_window(self):
        """Place the window at the bottom-right of the primary screen."""
        sw = self._win.winfo_screenwidth()
        sh = self._win.winfo_screenheight()
        x  = sw - WIDGET_W - MARGIN_RIGHT
        y  = sh - WIDGET_H - MARGIN_BOTTOM
        self._win.geometry(f"{WIDGET_W}x{WIDGET_H}+{x}+{y}")

    # ------------------------------------------------------------------
    # Drag / reposition
    # ------------------------------------------------------------------

    def _drag_start(self, event):
        self._drag_x = event.x
        self._drag_y = event.y

    def _drag_motion(self, event):
        dx = event.x - self._drag_x
        dy = event.y - self._drag_y
        x  = self._win.winfo_x() + dx
        y  = self._win.winfo_y() + dy
        self._win.geometry(f"+{x}+{y}")

    # ------------------------------------------------------------------
    # Click handlers
    # ------------------------------------------------------------------

    def _on_left_click(self, event):
        if self._on_click:
            self._on_click()

    def _show_context_menu(self, event):
        try:
            self._context_menu.tk_popup(event.x_root, event.y_root)
        finally:
            self._context_menu.grab_release()

    def _open_detailed(self):
        if self._on_click:
            self._on_click()

    def _reset_session(self):
        if messagebox.askyesno(
            "Reset All Data",
            "Clear all recorded token data and rescan from scratch?\nThis cannot be undone.",
            parent=self._win
        ):
            self._storage.clear_all()
            # Also clear watcher offsets if a reset callback is set
            if hasattr(self, '_on_reset') and self._on_reset:
                self._on_reset()

    def _quit_app(self):
        self._running = False
        self._root.quit()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def update_display(self, data_dict):
        """
        Update the visible token counts.

        Expected keys in data_dict:
            input_tokens, output_tokens, cache_creation, cache_read,
            model, session_id, request_count, today_cost
        """
        inp        = data_dict.get("input_tokens",  0)
        out        = data_dict.get("output_tokens", 0)
        sid        = data_dict.get("session_id",    "—") or "—"
        today_cost = data_dict.get("today_cost",    "$0.00")

        # Truncate very long session IDs
        if len(str(sid)) > 20:
            sid = str(sid)[:18] + "..."

        self._label_counts.configure(
            text=f"In: {_fmt(inp)}  |  Out: {_fmt(out)}"
        )
        self._label_cost.configure(
            text=f"Today: {today_cost}"
        )
        self._label_session.configure(
            text=sid
        )

    def set_on_reset(self, callback):
        """Register a callback to be invoked when the user resets all data."""
        self._on_reset = callback

    def set_status(self, active: bool):
        """Set the status dot: green when proxy is active, gray otherwise."""
        self._dot.configure(fg=DOT_ACTIVE if active else DOT_IDLE)

    # ------------------------------------------------------------------
    # Auto-refresh loop
    # ------------------------------------------------------------------

    def _schedule_refresh(self):
        if not self._running:
            return
        if self._refresh_cb:
            try:
                self._refresh_cb()
            except Exception:
                pass
        self._win.after(2000, self._schedule_refresh)

    # ------------------------------------------------------------------
    # Lifecycle helpers
    # ------------------------------------------------------------------

    def destroy(self):
        """Cleanly destroy the overlay window."""
        self._running = False
        try:
            self._win.destroy()
        except tk.TclError:
            pass