import tkinter as tk
from tkinter import simpledialog, messagebox
import json
import os
from pynput import keyboard
import threading

# 配置文件路径
CONFIG_FILE = "key_config.json"

class KeyOverlayApp:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("KeyOverlay Py")
        
        # 初始配置
        self.config = {
            "alpha": 0.8,
            "position": "+100+100",
            "monitored_keys": ["w", "a", "s", "d", "space"],
            "counts": {}
        }
        
        self.load_config()
        
        # UI 设置
        self.root.overrideredirect(True)  # 无边框
        self.root.attributes("-topmost", True)  # 置顶
        self.root.attributes("-alpha", self.config["alpha"])
        self.root.geometry(self.config["position"])
        self.root.configure(bg='#282c34')

        # 标签字典，用于动态更新界面
        self.labels = {}
        self.setup_ui()

        # 拖动功能变量
        self._drag_data = {"x": 0, "y": 0}
        self.root.bind("<Button-1>", self.start_drag)
        self.root.bind("<B1-Motion>", self.do_drag)
        
        # 右键菜单
        self.menu = tk.Menu(self.root, tearoff=0)
        self.menu.add_command(label="添加按键", command=self.add_key_dialog)
        self.menu.add_command(label="删除按键", command=self.remove_key_dialog)
        self.menu.add_command(label="设置透明度", command=self.set_alpha_dialog)
        self.menu.add_command(label="重置计数", command=self.reset_counts)
        self.menu.add_separator()
        self.menu.add_command(label="退出并保存", command=self.quit_app)
        
        self.root.bind("<Button-3>", self.show_menu)

        # 启动监听线程
        self.listener = keyboard.Listener(on_press=self.on_press)
        self.listener.start()

        self.root.mainloop()

    def load_config(self):
        if os.path.exists(CONFIG_FILE):
            with open(CONFIG_FILE, "r") as f:
                try:
                    saved_config = json.load(f)
                    self.config.update(saved_config)
                except:
                    pass
        # 初始化缺失的计数
        for key in self.config["monitored_keys"]:
            if key not in self.config["counts"]:
                self.config["counts"][key] = 0

    def save_config(self):
        # 保存位置信息
        self.config["position"] = f"+{self.root.winfo_x()}+{self.root.winfo_y()}"
        with open(CONFIG_FILE, "w") as f:
            json.dump(self.config, f, indent=4)

    def setup_ui(self):
        # 清除旧标签
        for label in self.labels.values():
            label.destroy()
        self.labels = {}

        # 创建新标签
        for i, key in enumerate(self.config["monitored_keys"]):
            count = self.config["counts"].get(key, 0)
            lbl = tk.Label(self.root, text=f"{key.upper()}: {count}", 
                           fg="white", bg="#282c34", 
                           font=("Consolas", 12, "bold"),
                           padx=10, pady=5)
            lbl.pack(fill=tk.X)
            self.labels[key] = lbl

    def on_press(self, key):
        try:
            # 尝试获取按键名称
            if hasattr(key, 'char') and key.char is not None:
                k = key.char.lower()
            else:
                k = key.name.lower()
        except:
            return

        if k in self.config["monitored_keys"]:
            self.config["counts"][k] = self.config["counts"].get(k, 0) + 1
            # 在主线程更新 UI
            count = self.config["counts"][k]
            self.root.after(0, lambda: self.labels[k].config(text=f"{k.upper()}: {count}"))

    # 拖动逻辑
    def start_drag(self, event):
        self._drag_data["x"] = event.x
        self._drag_data["y"] = event.y

    def do_drag(self, event):
        x = self.root.winfo_x() - self._drag_data["x"] + event.x
        y = self.root.winfo_y() - self._drag_data["y"] + event.y
        self.root.geometry(f"+{x}+{y}")

    # 菜单功能
    def show_menu(self, event):
        self.menu.post(event.x_root, event.y_root)

    def add_key_dialog(self):
        new_key = simpledialog.askstring("添加按键", "输入按键名称 (如: q, space, shift):")
        if new_key:
            new_key = new_key.lower()
            if new_key not in self.config["monitored_keys"]:
                self.config["monitored_keys"].append(new_key)
                if new_key not in self.config["counts"]:
                    self.config["counts"][new_key] = 0
                self.setup_ui()

    def remove_key_dialog(self):
        key_to_remove = simpledialog.askstring("删除按键", "输入要删除的按键名称:")
        if key_to_remove and key_to_remove.lower() in self.config["monitored_keys"]:
            self.config["monitored_keys"].remove(key_to_remove.lower())
            self.setup_ui()

    def set_alpha_dialog(self):
        new_alpha = simpledialog.askfloat("设置透明度", "输入透明度 (0.1 - 1.0):", minvalue=0.1, maxvalue=1.0)
        if new_alpha:
            self.config["alpha"] = new_alpha
            self.root.attributes("-alpha", new_alpha)

    def reset_counts(self):
        if messagebox.askyesno("重置", "确定要清空所有按键计数吗？"):
            for key in self.config["counts"]:
                self.config["counts"][key] = 0
            self.setup_ui()

    def quit_app(self):
        self.save_config()
        self.root.destroy()
        os._exit(0)

if __name__ == "__main__":
    KeyOverlayApp()