mempref_gui.py

import tkinter as tk
from tkinter import messagebox
import ctypes, sys, time
from collections import deque

from mempref_cli import WinRing0, ASMMAP, IMCEvent, bytes_to_human_readable

class MemPerfGUI:
    """内存性能监控悬浮窗"""

    def __init__(self):
        self.root = tk.Tk()
        self.root.title("MemPerfMon")
        self.root.overrideredirect(True)
        self.root.wm_attributes("-topmost", True)
        self.root.wm_attributes("-alpha", 0.88)
        self.root.configure(bg="#0d1117")

        # 定位到屏幕右上角
        sw = self.root.winfo_screenwidth()
        self.root.geometry(f"240x120+{sw - 270}+40")

        self._drag_x = 0
        self._drag_y = 0
        self.drivers_ok = False
        self.chart_visible = False
        self.history = deque(maxlen=100)  # [(read_bw, write_bw, total_bw), ...]

        self._build_ui()
        self._init_drivers()
        self._start_update()

    def _build_ui(self):
        # ---- 标题栏 ----
        title_bar = tk.Frame(self.root, bg="#161b22", height=26, cursor="fleur")
        title_bar.pack(fill=tk.X)
        title_bar.bind("<Button-1>", self._drag_start)
        title_bar.bind("<B1-Motion>", self._drag_move)

        title_label = tk.Label(
            title_bar,
            text=" MemPerfMon",
            bg="#161b22",
            fg="#58a6ff",
            font=("Segoe UI", 10, "bold"),
        )
        title_label.pack(side=tk.LEFT, padx=6)
        title_label.bind("<Button-1>", self._drag_start)
        title_label.bind("<B1-Motion>", self._drag_move)

        close_btn = tk.Label(
            title_bar,
            text="✕",
            bg="#161b22",
            fg="#8b949e",
            font=("Segoe UI", 11),
            cursor="hand2",
        )
        close_btn.pack(side=tk.RIGHT, padx=8)
        close_btn.bind("<Button-1>", lambda e: self._shutdown())
        close_btn.bind("<Enter>", lambda e: close_btn.config(fg="#f85149"))
        close_btn.bind("<Leave>", lambda e: close_btn.config(fg="#8b949e"))

        # ---- 外层容器（避免 pack 顺序在动态 pack_forget 后错乱） ----
        self.outer = tk.Frame(self.root, bg="#0d1117")
        self.outer.pack(fill=tk.BOTH, expand=True)

        # ---- 主区域（水平容器：左侧数据 + 右侧图表） ----
        self.main_area = tk.Frame(self.outer, bg="#0d1117")
        self.main_area.pack(fill=tk.BOTH, expand=True, padx=12, pady=(6, 4))

        # 左侧列：实时带宽 + 累计统计
        self.left_panel = tk.Frame(self.main_area, bg="#0d1117")
        self.left_panel.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 数据区
        data = tk.Frame(self.left_panel, bg="#0d1117")
        data.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
        self.data_frame = data

        self.labels = {}
        colors = {"Read": "#3fb950",
                  "Write": "#d29922",
                  "Total": "#f78166",
                  }
        for i, (key, color) in enumerate(colors.items()):
            tk.Label(
                data,
                text=f"{key}:",
                bg="#0d1117",
                fg="#8b949e",
                font=("Consolas", 11),
                width=6,
                anchor=tk.W,
            ).grid(row=i, column=0, sticky=tk.W)

            val = tk.Label(
                data,
                text="--.-- /s",
                bg="#0d1117",
                fg=color,
                font=("Consolas", 11, "bold"),
                anchor=tk.E,
            )
            val.grid(row=i, column=1, sticky=tk.E, padx=(4, 0))
            self.labels[key] = val

        data.columnconfigure(1, weight=1)

        # 累计统计区（位于实时带宽下方）
        self.cumulative_visible = False
        cum = tk.Frame(self.left_panel, bg="#0d1117")

        self.cumulative_labels = {}
        cum_keys = {"Total Read": "#3fb950",
                     "Total Write": "#d29922",
                     "Total Access": "#f78166",
                     }
        for i, (key, color) in enumerate(cum_keys.items()):
            tk.Label(
                cum,
                text=f"{key}:",
                bg="#0d1117",
                fg="#8b949e",
                font=("Consolas", 11),
                width=12,
                anchor=tk.W,
            ).grid(row=i, column=0, sticky=tk.W)

            val = tk.Label(
                cum,
                text="--.--",
                bg="#0d1117",
                fg=color,
                font=("Consolas", 11, "bold"),
                anchor=tk.E,
            )
            val.grid(row=i, column=1, sticky=tk.E, padx=(4, 0))
            self.cumulative_labels[key] = val

        cum.columnconfigure(1, weight=1)
        self.cumulative_frame = cum

        # 图表区（初始隐藏）
        self.chart_canvas = tk.Canvas(
            self.main_area,
            bg="#161b22",
            highlightthickness=0,
        )

        # 右键菜单
        self.menu = tk.Menu(self.root, tearoff=0, bg="#161b22", fg="#c9d1d9")
        self.menu.add_command(label="累计统计", command=self._show_cumulative_stats)
        self.menu.add_command(label="历史带宽折线图", command=self._toggle_chart)
        self.menu.add_separator()
        self.menu.add_command(label="退出", command=self._shutdown)
        self.root.bind("<Button-3>", lambda e: self.menu.post(e.x_root, e.y_root))

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

    def _drag_move(self, event):
        x = self.root.winfo_x() + event.x - self._drag_x
        y = self.root.winfo_y() + event.y - self._drag_y
        self.root.geometry(f"+{x}+{y}")

    def _init_drivers(self):
        try:
            self.winring0 = WinRing0()
            self.asmmap = ASMMAP()
            cpu_model = self.winring0.getCPUFamilyModelFromCPUID(self.winring0.cpuid(1))
            if cpu_model[0] != 6:
                raise RuntimeError(f"Unsupported CPU family: {cpu_model[0]}")
            mchbar = self.winring0.ReadPciConfigDword(0, 0, 0, 0x48) & 0xFFFFF000
            if mchbar == 0:
                raise RuntimeError("不支持当前平台")
            self.imc_event = IMCEvent(cpu_model, mchbar, self.asmmap)
            self.prev_total, self.prev_read, self.prev_write = self.imc_event.get_metrics()
            self.last_time = time.time()
            self.drivers_ok = True
        except Exception as e:
            messagebox.showerror("初始化失败", str(e))
            self.root.destroy()
            sys.exit(1)

    def _start_update(self):
        self.root.after(500, self._update_loop)
        self.root.mainloop()

    def _update_loop(self):
        if not self.drivers_ok:
            return
        try:
            cur_total, cur_read, cur_write = self.imc_event.get_metrics()
            now = time.time()
            dt = now - self.last_time
            if dt > 0:
                read_bw = (cur_read - self.prev_read) * self.imc_event.IMC_ACCESS_SIZE / dt
                write_bw = (cur_write - self.prev_write) * self.imc_event.IMC_ACCESS_SIZE / dt
                total_bw = (cur_total - self.prev_total) * self.imc_event.IMC_ACCESS_SIZE / dt

                self.labels["Read"].config(text=f"{bytes_to_human_readable(read_bw)}/s")
                self.labels["Write"].config(
                    text=f"{bytes_to_human_readable(write_bw)}/s"
                )
                self.labels["Total"].config(
                    text=f"{bytes_to_human_readable(total_bw)}/s"
                )

            self.prev_total, self.prev_read, self.prev_write = (
                cur_total,
                cur_read,
                cur_write,
            )
            self.last_time = now

            # 历史数据记录：始终保留最近 200 个点，图表仅负责展示
            self.history.append((read_bw, write_bw, total_bw))

            # 折线图刷新
            if self.chart_visible:
                self._draw_chart()

            # 同步刷新累计统计区
            if self.cumulative_visible:
                total_access = cur_total * self.imc_event.IMC_ACCESS_SIZE
                total_read = cur_read * self.imc_event.IMC_ACCESS_SIZE
                total_write = cur_write * self.imc_event.IMC_ACCESS_SIZE
                self.cumulative_labels["Total Read"].config(
                    text=bytes_to_human_readable(total_read)
                )
                self.cumulative_labels["Total Write"].config(
                    text=bytes_to_human_readable(total_write)
                )
                self.cumulative_labels["Total Access"].config(
                    text=bytes_to_human_readable(total_access)
                )
        except Exception:
            pass

        self.root.after(500, self._update_loop)

    def _update_window_geometry(self):
        """根据当前可见面板状态同步窗口大小。"""
        if self.chart_visible and self.cumulative_visible:
            width, height = 500, 220
        elif self.chart_visible:
            width, height = 500, 120
        elif self.cumulative_visible:
            width, height = 240, 220
        else:
            width, height = 240, 120

        x = self.root.winfo_x()
        y = self.root.winfo_y()
        self.root.geometry(f"{width}x{height}+{x}+{y}")

    def _show_cumulative_stats(self):
        """切换累计统计区的显示/隐藏"""
        if self.cumulative_visible:
            self.cumulative_frame.pack_forget()
            self.cumulative_visible = False
        else:
            self.cumulative_frame.pack(side=tk.TOP, fill=tk.X, padx=(0, 12), pady=(4, 6))
            self.cumulative_visible = True

        self._update_window_geometry()

    def _toggle_chart(self):
        """切换历史带宽折线图的显示/隐藏"""
        if self.chart_visible:
            self.chart_canvas.pack_forget()
            self.chart_visible = False
        else:
            self.chart_canvas.pack(
                side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=(8, 0)
            )
            self.chart_visible = True

        self._update_window_geometry()

    def _draw_chart(self):
        """在 canvas 上绘制三条历史带宽折线"""
        c = self.chart_canvas
        c.delete("all")
        w = c.winfo_width()
        h = c.winfo_height()
        if w < 10 or h < 10 or len(self.history) < 2:
            return

        margin = 6
        plot_left = margin + 40   # 留出 Y 轴标签空间
        plot_top = margin
        plot_w = w - plot_left - margin
        plot_h = h - 2 * margin

        # 计算最大值用于 Y 轴缩放
        max_val = max(max(r, wb, t) for r, wb, t in self.history)
        if max_val == 0:
            max_val = 1

        def scale_x(i):
            return plot_left + (i / max(len(self.history) - 1, 1)) * plot_w

        def scale_y(v):
            return plot_top + plot_h - (v / max_val) * plot_h

        # 背景网格 (水平线)
        for level in range(5):
            y = plot_top + (level / 4) * plot_h
            c.create_line(plot_left, y, plot_left + plot_w, y,
                          fill="#30363d", dash=(2, 4))

        # Y 轴标签
        for level in range(5):
            val = max_val * (1 - level / 4)
            y = plot_top + (level / 4) * plot_h
            c.create_text(plot_left - 4, y,
                          text=self._fmt_bw(val), anchor=tk.E,
                          fill="#8b949e", font=("Consolas", 7))

        # 绘制三条折线
        lines = [
            ("Read",  "#3fb950", 0),
            ("Write", "#d29922", 1),
            ("Total", "#f78166", 2),
        ]
        for _name, color, idx in lines:
            if len(self.history) < 2:
                continue
            points = []
            for i, hdata in enumerate(self.history):
                x = scale_x(i)
                y = scale_y(hdata[idx])
                points.extend((x, y))
            if points:
                c.create_line(*points, fill=color, width=1.5, smooth=False)

        # 图例
        legend_x = plot_left + 4
        legend_y = plot_top + 2
        for name, color, _idx in lines:
            c.create_line(legend_x, legend_y + 5, legend_x + 12, legend_y + 5,
                          fill=color, width=1.5)
            c.create_text(legend_x + 16, legend_y + 5,
                          text=name, anchor=tk.W,
                          fill=color, font=("Consolas", 7))
            legend_y += 12

    @staticmethod
    def _fmt_bw(val):
        """格式化带宽值用于 Y 轴标签（简洁版）"""
        if val >= 1e12:
            return f"{val/1e12:.0f}TB"
        elif val >= 1e9:
            return f"{val/1e9:.0f}GB"
        elif val >= 1e6:
            return f"{val/1e6:.0f}MB"
        elif val >= 1e3:
            return f"{val/1e3:.0f}KB"
        else:
            return f"{val:.0f}B"

    def _shutdown(self):
        self.drivers_ok = False
        self.root.destroy()


def run_as_admin():
    """以管理员权限重新启动自身"""
    if not ctypes.windll.shell32.IsUserAnAdmin():
        ctypes.windll.shell32.ShellExecuteW(
            None, "runas", sys.executable, " ".join(sys.argv), None, 1
        )
        return True
    return False


if __name__ == "__main__":
    if run_as_admin():
        sys.exit()
    MemPerfGUI()