tray.py

"""
MousePal — System Tray
"""

import threading
from pathlib import Path
from PIL import Image, ImageDraw
import pystray

PROFILE_COLOURS = {
    "Precision": (80, 140, 255),    # Blue
    "Everyday": (50, 200, 80),      # Green
    "Gaming": (220, 50, 50),        # Red
}
DEFAULT_COLOUR = (100, 180, 255)
LOCK_COLOUR = (200, 50, 50)

_logo_cache = {}


def _load_logo(size):
    if size in _logo_cache:
        return _logo_cache[size]
    logo_path = Path(__file__).parent / "assets" / "mousewheel_logo.png"
    if logo_path.exists():
        logo = Image.open(str(logo_path)).convert("RGBA")
        logo_size = int(size * 0.6)
        logo = logo.resize((logo_size, logo_size), Image.LANCZOS)
        _logo_cache[size] = logo
        return logo
    return None


def _generate_icon(colour, size=64, locked=False):
    img = Image.new("RGBA", (size, size), (0, 0, 0, 0))
    draw = ImageDraw.Draw(img)

    margin = size // 16
    ring_width = max(3, size // 10)
    draw.ellipse([margin, margin, size - margin, size - margin],
                 fill=colour + (255,))

    inner_margin = margin + ring_width
    draw.ellipse([inner_margin, inner_margin,
                  size - inner_margin, size - inner_margin],
                 fill=(20, 25, 35, 255))

    logo = _load_logo(size)
    if logo:
        lx = (size - logo.width) // 2
        ly = (size - logo.height) // 2
        img.paste(logo, (lx, ly), logo)

    if locked:
        lr = size * 0.14
        lx = size - margin - lr * 1.5
        ly = size - margin - lr * 1.5
        draw.ellipse([lx, ly, lx + lr * 2, ly + lr * 2],
                     fill=LOCK_COLOUR + (255,),
                     outline=(255, 255, 255, 200), width=max(1, size // 32))

    return img


class TrayApp:
    def __init__(self, profile_manager, config, on_settings=None, on_quit=None,
                 on_find_cursor=None, on_magnifier=None):
        self.pm = profile_manager
        self.config = config
        self._on_settings = on_settings
        self._on_quit = on_quit
        self._on_find_cursor = on_find_cursor
        self._on_magnifier = on_magnifier
        self._icon = None
        self._thread = None

        self._icons = {}
        self._icons_locked = {}
        self._default_icon = _generate_icon(DEFAULT_COLOUR)
        self._rebuild_icons()

    def _rebuild_icons(self):
        for name in self.pm.get_profile_names():
            if name not in self._icons:
                colour = PROFILE_COLOURS.get(name, DEFAULT_COLOUR)
                self._icons[name] = _generate_icon(colour)
                self._icons_locked[name] = _generate_icon(colour, locked=True)

    def refresh_profiles(self):
        self._rebuild_icons()
        self._update_icon()

    def _get_icon_for_profile(self, name):
        if self.pm.is_locked():
            return self._icons_locked.get(name, self._default_icon)
        return self._icons.get(name, self._default_icon)

    def _build_menu(self):
        return pystray.Menu(lambda: self._build_menu_items())

    def _build_menu_items(self):
        items = []
        for name in self.pm.get_profile_names():
            items.append(
                pystray.MenuItem(
                    name,
                    self._make_switch_handler(name),
                    checked=lambda item, n=name: self.pm.get_active() == n,
                    radio=True,
                )
            )

        # Mouse Speed submenu
        items.append(pystray.MenuItem(
            "Mouse Speed",
            pystray.Menu(
                pystray.MenuItem("+2 (coarse)", lambda i, it: self._nudge_speed(2)),
                pystray.MenuItem("+1 (fine)",   lambda i, it: self._nudge_speed(1)),
                pystray.MenuItem("-1 (fine)",   lambda i, it: self._nudge_speed(-1)),
                pystray.MenuItem("-2 (coarse)", lambda i, it: self._nudge_speed(-2)),
            )
        ))

        # Scroll Speed submenu
        items.append(pystray.MenuItem(
            "Scroll Speed",
            pystray.Menu(
                pystray.MenuItem("+5 lines", lambda i, it: self._nudge_scroll(5)),
                pystray.MenuItem("+1 line",  lambda i, it: self._nudge_scroll(1)),
                pystray.MenuItem("-1 line",  lambda i, it: self._nudge_scroll(-1)),
                pystray.MenuItem("-5 lines", lambda i, it: self._nudge_scroll(-5)),
            )
        ))

        # Cursor Size submenu
        items.append(pystray.MenuItem(
            "Cursor Size",
            pystray.Menu(
                pystray.MenuItem("Larger",  lambda i, it: self._nudge_cursor(1)),
                pystray.MenuItem("Smaller", lambda i, it: self._nudge_cursor(-1)),
            )
        ))

        lock_text = "Unlock Profile" if self.pm.is_locked() else "Lock Profile"

        items.extend([
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("Find My Cursor", self._on_find_cursor_click),
            pystray.MenuItem("Magnifier", self._on_magnifier_click),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("PANIC! (Reset defaults)", self._on_panic_click),
            pystray.MenuItem(lock_text, self._on_lock_click),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("Settings", self._on_settings_click),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("Quit", self._on_quit_click),
        ])
        return items

    def _make_switch_handler(self, profile_name):
        def handler(icon, item):
            self.pm.switch(profile_name, force=True)
        return handler

    def _nudge_speed(self, delta):
        import mouse
        new_value = mouse.nudge_mouse_speed(delta)
        if self._icon:
            active = self.pm.get_active()
            self._icon.title = f"MousePal - {active} | Speed {new_value}"
            try:
                self._icon.update_menu()
            except Exception:
                pass

    def _nudge_scroll(self, delta):
        import mouse
        new_value = mouse.nudge_scroll_lines(delta)
        if self._icon:
            active = self.pm.get_active()
            self._icon.title = f"MousePal - {active} | Scroll {new_value}"
            try:
                self._icon.update_menu()
            except Exception:
                pass

    def _nudge_cursor(self, delta):
        import mouse
        new_value = mouse.nudge_cursor_size(delta)
        if self._icon:
            active = self.pm.get_active()
            self._icon.title = f"MousePal - {active} | Cursor size {new_value}"
            try:
                self._icon.update_menu()
            except Exception:
                pass

    def _on_find_cursor_click(self, icon, item):
        if self._on_find_cursor:
            self._on_find_cursor()

    def _on_magnifier_click(self, icon, item):
        if self._on_magnifier:
            self._on_magnifier()

    def _on_panic_click(self, icon, item):
        self.pm.switch("Everyday", force=True)

    def _on_lock_click(self, icon, item):
        self.pm.toggle_lock()
        self._update_icon()

    def _on_settings_click(self, icon, item):
        if self._on_settings:
            self._on_settings()

    def _on_quit_click(self, icon, item):
        if self._on_quit:
            self._on_quit()

    def start(self):
        active = self.pm.get_active()
        self._icon = pystray.Icon(
            "MousePal",
            self._get_icon_for_profile(active),
            title=self._build_title(active),
            menu=self._build_menu(),
        )
        self._thread = threading.Thread(target=self._icon.run, daemon=True)
        self._thread.start()

    def stop(self):
        if self._icon:
            try:
                self._icon.stop()
            except Exception:
                pass

    def _build_title(self, profile_name):
        locked = " [LOCKED]" if self.pm.is_locked() else ""
        return f"MousePal - {profile_name}{locked}"

    def update_tooltip(self, profile_name):
        if self._icon:
            self._icon.title = self._build_title(profile_name)
            self._icon.icon = self._get_icon_for_profile(profile_name)
            try:
                self._icon.update_menu()
            except Exception:
                pass
            if self.config.get("notifications_enabled", True):
                try:
                    self._icon.notify(f"Switched to {profile_name}", "MousePal")
                except Exception:
                    pass

    def _update_icon(self):
        if self._icon:
            active = self.pm.get_active()
            self._icon.title = self._build_title(active)
            self._icon.icon = self._get_icon_for_profile(active)
            try:
                self._icon.update_menu()
            except Exception:
                pass