proxy-god/proxy_chain_manager/tray.py

"""System tray: LED-style proxy on/off + hover tooltip with exit IP."""
from __future__ import annotations

import logging
import threading
from typing import Any, Callable

log = logging.getLogger(__name__)

import pystray
from PIL import Image, ImageDraw, ImageFilter

_SIZE = 64


def _hex_rgb(h: str) -> tuple[int, int, int]:
    h = h.lstrip("#")
    return int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)


def _led_icon(led: str, glow: str, rim: str = "#2d3748") -> Image.Image:
    """Traffic-light style icon: dark housing + glowing LED."""
    base = Image.new("RGBA", (_SIZE, _SIZE), (0, 0, 0, 0))
    draw = ImageDraw.Draw(base)
    cx, cy = _SIZE // 2, _SIZE // 2

    # Outer housing (rounded rect feel via ellipse)
    draw.ellipse([4, 4, _SIZE - 4, _SIZE - 4], fill="#0d0d1a", outline=rim, width=2)
    draw.ellipse([10, 10, _SIZE - 10, _SIZE - 10], fill="#12122a", outline="#1a1a3e", width=1)

    glow_layer = Image.new("RGBA", (_SIZE, _SIZE), (0, 0, 0, 0))
    g = ImageDraw.Draw(glow_layer)
    lr, lg, lb = _hex_rgb(glow)
    for radius, alpha in ((22, 35), (16, 70), (11, 120)):
        g.ellipse(
            [cx - radius, cy - radius, cx + radius, cy + radius],
            fill=(lr, lg, lb, alpha),
        )
    glow_layer = glow_layer.filter(ImageFilter.GaussianBlur(radius=2))
    base = Image.alpha_composite(base, glow_layer)

    draw = ImageDraw.Draw(base)
    dr, dg, db = _hex_rgb(led)
    draw.ellipse([cx - 9, cy - 9, cx + 9, cy + 9], fill=(dr, dg, db, 255))
    draw.ellipse([cx - 5, cy - 6, cx + 2, cy - 1], fill=(255, 255, 255, 90))

    return base


ICONS = {
    "green": _led_icon("#00e676", "#00e676"),
    "red": _led_icon("#ff1744", "#ff1744"),
    "yellow": _led_icon("#ffab00", "#ffab00"),
    "gray": _led_icon("#4a5568", "#3d4a5c", rim="#374151"),
}

_STATE_LABEL = {
    "green": "PROXY ON",
    "red": "PROXY OFF / ERROR",
    "yellow": "CONNECTING…",
    "gray": "STOPPED",
}


def _tooltip(state: str, exit_ip: str | None) -> str:
    label = _STATE_LABEL.get(state, "Proxy God")
    if exit_ip:
        return f"Proxy God — {label}\nExit IP: {exit_ip}"
    if state == "green":
        return f"Proxy God — {label}\nExit IP: (checking…)"
    return f"Proxy God — {label}\nExit IP: —"


class TrayIcon:
    def __init__(
        self,
        on_show: Callable[[], None],
        on_quit: Callable[[], None],
        on_rotate: Callable[[], None],
    ) -> None:
        self._on_show = on_show
        self._on_quit = on_quit
        self._on_rotate = on_rotate
        self._icon: pystray.Icon | None = None
        self._thread: threading.Thread | None = None
        self._state = "gray"
        self._exit_ip: str | None = None
        self._lock = threading.Lock()

    def start(self) -> None:
        if self._thread and self._thread.is_alive():
            return
        menu = pystray.Menu(
            pystray.MenuItem("Show Proxy God", self._show, default=True),
            pystray.MenuItem("Rotate chain now", self._rotate),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("Quit", self._quit),
        )
        with self._lock:
            tip = _tooltip(self._state, self._exit_ip)
            self._icon = pystray.Icon(
                "ProxyGod",
                icon=ICONS[self._state],
                title=tip,
                menu=menu,
            )
        self._thread = threading.Thread(target=self._icon.run, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        if self._icon:
            try:
                self._icon.stop()
            except Exception as exc:
                log.warning("Tray icon stop failed: %s", exc)

    def set_state(self, state: str, exit_ip: str | None = None) -> None:
        """state: green (on), red (error), yellow (connecting), gray (stopped).

        Pass exit_ip to refresh the hover tooltip. Use exit_ip=None to keep the last IP.
        """
        if state not in ICONS:
            state = "gray"
        with self._lock:
            self._state = state
            if exit_ip is not None:
                self._exit_ip = exit_ip.strip() if exit_ip else None
            if self._icon:
                self._icon.icon = ICONS[state]
                self._icon.title = _tooltip(state, self._exit_ip)

    def set_exit_ip(self, exit_ip: str | None) -> None:
        """Update tooltip only (e.g. after health check, same LED color)."""
        with self._lock:
            self._exit_ip = exit_ip.strip() if exit_ip else None
            if self._icon:
                self._icon.title = _tooltip(self._state, self._exit_ip)

    def _show(self, icon: Any = None, item: Any = None) -> None:
        self._on_show()

    def _rotate(self, icon: Any = None, item: Any = None) -> None:
        self._on_rotate()

    def _quit(self, icon: Any = None, item: Any = None) -> None:
        self._on_quit()