tray_icon.py

# tray_icon.py

import pystray
from PIL import Image, ImageDraw
import threading
import os
import sys
import time
from gui import run_gui 
from modeswitch import AppMode # <--- 新增导入

class TrayIcon:
    def __init__(self, mode_switch, listener, stop_event):
        self.mode_switch = mode_switch
        self.listener = listener
        self.stop_event = stop_event
        self.icon = None
        self.active_icon = self.create_icon("active")
        self.inactive_icon = self.create_icon("inactive")
        self.command_mode_icon = self.create_icon("command")
        self.key_mapping_mode_icon = self.create_icon("key_mapping")
        self.thread = None
        self.is_settings_window_open = False
        self.is_exiting = False
        self.cleanup_completed = False  # 标志位：表示是否已完成清理

    def force_cleanup(self):
        """强制清理所有托盘图标资源"""
        try:
            # 设置清理完成标志
            self.cleanup_completed = True

            # 停止图标事件循环
            if self.icon is not None:
                try:
                    self.icon.stop()
                except Exception as e:
                    pass  # 忽略停止过程中的错误

            # 清理图标引用
            self.icon = None

            # 如果线程还在运行，等待其结束
            if self.thread is not None and self.thread.is_alive():
                try:
                    self.thread.join(timeout=1.0)
                except Exception as e:
                    pass  # 忽略等待过程中的错误

            # 清理线程引用
            self.thread = None

        except Exception as e:
            pass  # 忽略所有清理过程中的错误

    def run(self):
        """运行托盘图标 - 在运行前检查是否需要清理"""
        # 如果之前有未完成的清理，先强制清理
        if hasattr(self, 'cleanup_completed') and not self.cleanup_completed:
            self.force_cleanup()

        self.thread = threading.Thread(target=self._run, daemon=True)
        self.thread.start()

    def create_icon(self, state):
        width, height = 64, 64
        image = Image.new('RGB', (width, height), color=(0, 0, 0, 0))
        dc = ImageDraw.Draw(image)
        if state == "active":
            dc.rectangle([(8, 8), (width-8, height-8)], fill=(0, 180, 0))
            dc.ellipse([(16, 16), (width-16, height-16)], fill=(0, 220, 0))
        elif state == "inactive":
            dc.rectangle([(8, 8), (width-8, height-8)], fill=(100, 100, 100))
            dc.ellipse([(16, 16), (width-16, height-16)], fill=(150, 150, 150))
        elif state == "command":
            dc.rectangle([(8, 8), (width-8, height-8)], fill=(0, 0, 180))
            dc.ellipse([(16, 16), (width-16, height-16)], fill=(0, 0, 220))
        elif state == "key_mapping":
            dc.rectangle([(8, 8), (width-8, height-8)], fill=(180, 0, 0))
            dc.ellipse([(16, 16), (width-16, height-16)], fill=(220, 0, 0))
        return image
    
    def update_icon(self):
        """根据当前的程序模式，更新托盘图标和标题。"""
        if self.icon:
            if self.mode_switch.is_mouse_control_mode():
                self.icon.icon = self.active_icon
            elif self.mode_switch.is_command_mode():
                self.icon.icon = self.command_mode_icon
            elif self.mode_switch.current_mode == AppMode.KEY_MAPPING:
                self.icon.icon = self.key_mapping_mode_icon
            else:
                self.icon.icon = self.inactive_icon
            self.icon.title = f"KeyMouse - {self.mode_switch.current_mode.name} 模式"
    
    def on_exit(self):
        """执行异步的、健壮的退出序列。"""
        if not self.is_exiting:
            self.is_exiting = True
            print("收到退出请求，启动后台关闭线程...")
            shutdown_thread = threading.Thread(target=self._shutdown_sequence, daemon=True)
            shutdown_thread.start()

    def _shutdown_sequence(self):
        """在后台执行所有实际的关闭操作，以避免死锁。"""
        import logging
        logging.info(">>> TrayIcon 后台关闭序列开始")
        print("后台关闭序列已启动...")

        try:
            if self.icon:
                logging.info("正在停止托盘图标...")
                self.icon.stop()
                logging.info("托盘图标已停止")
            else:
                logging.info("托盘图标不存在")
        except Exception as e:
            logging.error(f"停止托盘图标时发生异常: {e}", exc_info=True)

        try:
            logging.info("正在设置停止事件...")
            self.stop_event.set()
            logging.info("停止事件已设置")
        except Exception as e:
            logging.error(f"设置停止事件时发生异常: {e}", exc_info=True)

        try:
            if self.listener:
                logging.info("正在停止键盘监听器...")
                self.listener.stop()
                logging.info("键盘监听器停止信号已发送")

                if self.listener.is_alive():
                    logging.info("等待键盘监听器线程结束...")
                    self.listener.join(timeout=1.0)
                    if self.listener.is_alive():
                        logging.warning("键盘监听器线程在1秒内未结束")
                    else:
                        logging.info("键盘监听器线程已结束")
                else:
                    logging.info("键盘监听器线程已死亡")
            else:
                logging.info("键盘监听器不存在")
        except Exception as e:
            logging.error(f"停止键盘监听器时发生异常: {e}", exc_info=True)

        logging.info(">>> TrayIcon 后台关闭序列完成，准备强制退出")
        print("所有线程已停止。")

        os._exit(0)
    
    def on_settings_window_closed(self):
        self.is_settings_window_open = False
        print("设置窗口已关闭。")
    
    def on_settings(self):
        if self.is_settings_window_open:
            print("设置窗口已打开，请勿重复点击。")
            return
        
        print("正在打开设置窗口...")
        self.is_settings_window_open = True
        
        settings_thread = threading.Thread(
            target=run_gui, 
            args=(self.on_settings_window_closed, self),
            daemon=True
        )
        settings_thread.start()
    
    def run(self):
        self.thread = threading.Thread(target=self._run, daemon=True)
        self.thread.start()
    
    def _run(self):
        """在独立线程中运行pystray的事件循环。"""
        # 强制清理检查：如果之前有图标实例，先清理
        if self.icon is not None:
            try:
                self.icon.stop()
            except:
                pass
            self.icon = None

        # 创建新的图标实例
        menu = pystray.Menu(
            pystray.MenuItem("设置", self.on_settings),
            # DEL: [User Feedback] Removed "切换按键映射模式" as per user request.
            # pystray.MenuItem("切换按键映射模式", self.toggle_key_mapping_mode),
            pystray.MenuItem("退出", self.on_exit)
        )
        self.icon = pystray.Icon("keymouse", self.inactive_icon, "KeyMouse", menu)

        # --- 核心修正点：调用正确的新方法 ---
        # 1. 保存原始的模式切换方法
        original_toggle_method = self.mode_switch.toggle_mouse_control_mode

        # 2. 创建一个新的包装方法
        def new_toggle_with_icon_update():
            # 首先，调用原始的切换逻辑
            original_toggle_method()
            # 然后，更新我们的托盘图标
            self.update_icon()

        # 3. 用我们新的包装方法，替换掉原始的切换方法
        #    这样，每当主程序调用 toggle_mouse_control_mode 时，图标都会自动更新
        self.mode_switch.toggle_mouse_control_mode = new_toggle_with_icon_update

        # 首次运行时，也更新一次图标
        self.update_icon()

        # 运行图标事件循环
        if self.icon:
            self.icon.run()

        # 事件循环结束后，标记清理完成
        self.cleanup_completed = True

    # DEL: [User Feedback] Removed toggle_key_mapping_mode method as per user request.
    # def toggle_key_mapping_mode(self):
    #     """切换按键映射模式并更新图标。"""
    #     self.mode_switch.toggle_key_mapping_mode()
    #     self.update_icon()