pucoding_widget.py

"""Pucoding 用量桌面小组件（Windows / Python 3.10+）。

只调用 Pucoding 的只读 Key Portal 接口，不会把 Key 放进 URL，也不会发送给
Pucoding 以外的服务。Key 使用 Windows DPAPI 加密后保存在当前用户的 AppData 目录。
"""

from __future__ import annotations

import base64
import ctypes
import ctypes.wintypes as wintypes
import json
import os
import queue
import sys
import threading
import tkinter as tk
from datetime import datetime
from pathlib import Path
from tkinter import ttk
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


API_ROOT = "https://pucoding.com/api/v1"
SUMMARY_URL = f"{API_ROOT}/key-portal/summary"
USAGE_URL = f"{API_ROOT}/key-portal/usage/summary"
POLL_MS = 60_000
VERSION = "1.0.0"

APP_DIR = Path(os.environ.get("APPDATA", Path.home())) / "PucodingWidget"
CONFIG_PATH = APP_DIR / "config.json"


def protect_text(value: str) -> str:
    """Use Windows DPAPI so the saved key can only be decrypted by this user."""
    if os.name != "nt":
        raise OSError("API Key 加密仅支持 Windows")

    class DATA_BLOB(ctypes.Structure):
        _fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]

    raw = value.encode("utf-8")
    in_buf = ctypes.create_string_buffer(raw)
    in_blob = DATA_BLOB(len(raw), ctypes.cast(in_buf, ctypes.POINTER(ctypes.c_byte)))
    out_blob = DATA_BLOB()
    crypt32 = ctypes.windll.crypt32
    kernel32 = ctypes.windll.kernel32
    crypt32.CryptProtectData.restype = wintypes.BOOL
    kernel32.LocalFree.restype = wintypes.HLOCAL
    if not crypt32.CryptProtectData(ctypes.byref(in_blob), "Pucoding API Key", None, None, None, 0, ctypes.byref(out_blob)):
        raise OSError("Windows 无法保护 API Key")
    try:
        encrypted = ctypes.string_at(out_blob.pbData, out_blob.cbData)
    finally:
        kernel32.LocalFree(out_blob.pbData)
    return base64.b64encode(encrypted).decode("ascii")


def unprotect_text(value: str) -> str:
    if os.name != "nt":
        raise OSError("API Key 解密仅支持 Windows")

    class DATA_BLOB(ctypes.Structure):
        _fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]

    encrypted = base64.b64decode(value.encode("ascii"))
    in_buf = ctypes.create_string_buffer(encrypted)
    in_blob = DATA_BLOB(len(encrypted), ctypes.cast(in_buf, ctypes.POINTER(ctypes.c_byte)))
    out_blob = DATA_BLOB()
    crypt32 = ctypes.windll.crypt32
    kernel32 = ctypes.windll.kernel32
    crypt32.CryptUnprotectData.restype = wintypes.BOOL
    kernel32.LocalFree.restype = wintypes.HLOCAL
    if not crypt32.CryptUnprotectData(ctypes.byref(in_blob), None, None, None, None, 0, ctypes.byref(out_blob)):
        raise OSError("Windows 无法读取已保存的 API Key")
    try:
        plain = ctypes.string_at(out_blob.pbData, out_blob.cbData)
    finally:
        kernel32.LocalFree(out_blob.pbData)
    return plain.decode("utf-8")


def load_key() -> str:
    try:
        data = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
        encrypted = data.get("api_key", "")
        return unprotect_text(encrypted) if encrypted else ""
    except (OSError, ValueError, UnicodeError):
        return ""


def save_key(api_key: str) -> None:
    APP_DIR.mkdir(parents=True, exist_ok=True)
    payload = {"api_key": protect_text(api_key) if api_key else ""}
    temporary_path = CONFIG_PATH.with_suffix(".tmp")
    temporary_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
    temporary_path.replace(CONFIG_PATH)


def _number(value: object, default: float = 0.0) -> float:
    try:
        return float(value) if value is not None else default
    except (TypeError, ValueError):
        return default


def fetch_json(url: str, api_key: str) -> dict:
    request = Request(
        url,
        headers={
            "Accept": "application/json",
            "Cache-Control": "no-cache",
            "User-Agent": f"PucodingUsageWidget/{VERSION}",
            "x-api-key": api_key,
        },
    )
    try:
        with urlopen(request, timeout=20) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        if exc.code in (401, 403):
            raise RuntimeError("API Key 无效或已被禁用") from exc
        raise RuntimeError(f"服务器返回 HTTP {exc.code}") from exc
    except URLError as exc:
        raise RuntimeError("网络连接失败，请检查网络后重试") from exc
    except (TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc:
        raise RuntimeError("服务返回的数据无法解析") from exc

    if payload.get("code") not in (None, 200):
        raise RuntimeError(str(payload.get("message") or "查询失败"))
    data = payload.get("data", payload)
    if not isinstance(data, dict):
        raise RuntimeError("服务返回的数据格式不正确")
    return data


def read_usage(api_key: str) -> dict:
    """读取今日消费、累计消费和剩余额度。"""
    usage = fetch_json(USAGE_URL, api_key)
    summary = fetch_json(SUMMARY_URL, api_key)
    today = usage.get("today") if isinstance(usage.get("today"), dict) else {}
    key = usage.get("key") if isinstance(usage.get("key"), dict) else {}
    # summary 是备用来源，兼容接口字段随版本调整的情况。
    quota_used = _number(key.get("quota_used", summary.get("quota_used")))
    quota_limit = _number(key.get("quota_limit", summary.get("quota_limit")))
    remaining = key.get("quota_remaining", summary.get("quota_remaining"))
    return {
        "today_cost": _number(today.get("cost", summary.get("today_cost"))),
        "today_requests": int(_number(today.get("requests"))),
        "lifetime_cost": quota_used,
        "remaining": None if remaining is None else _number(remaining),
        "quota_limit": quota_limit,
        "request_count": int(_number(key.get("request_count", summary.get("request_count")))),
        "status": summary.get("status_name") or key.get("status_name") or "正常",
    }


class Widget(tk.Tk):
    BG = "#10131a"
    CARD = "#181d27"
    CARD_2 = "#202735"
    TEXT = "#f5f7fb"
    MUTED = "#9ca8bb"
    ACCENT = "#8de1c1"
    ORANGE = "#ffcc80"
    RED = "#ff8d9c"

    def __init__(self) -> None:
        super().__init__()
        self.title("Pucoding 用量")
        self.configure(bg=self.BG)
        self.geometry("280x190")
        self.minsize(240, 170)
        self.resizable(True, True)
        self._topmost = True
        self.attributes("-topmost", self._topmost)
        self.protocol("WM_DELETE_WINDOW", self._close)
        self.api_key = load_key()
        self._loading = False
        self._closing = False
        self._result_queue: queue.Queue[tuple[str, object, str]] = queue.Queue()
        self._poll_job: str | None = None
        self._auto_refresh_job: str | None = None
        self._build_style()
        self._build_ui()
        if self.api_key:
            self._set_status("正在读取…", self.MUTED)
            self.refresh()
        else:
            self._set_status("请先设置 API Key", self.ORANGE)
            self._show_settings()

    def _build_style(self) -> None:
        style = ttk.Style(self)
        style.theme_use("clam")
        style.configure("Widget.TButton", background=self.CARD_2, foreground=self.TEXT, borderwidth=0, padding=(10, 6), font=("Segoe UI", 9))
        style.map("Widget.TButton", background=[("active", "#2b3446")])
        style.configure("Accent.TButton", background=self.ACCENT, foreground="#0c1715", borderwidth=0, padding=(12, 7), font=("Segoe UI", 9, "bold"))
        style.map("Accent.TButton", background=[("active", "#a8efd5")])
        style.configure("Widget.TEntry", fieldbackground="#0d1016", foreground=self.TEXT, insertcolor=self.TEXT, borderwidth=0, padding=8)

    def _build_ui(self) -> None:
        self.columnconfigure(0, weight=1)
        self.rowconfigure(2, weight=1)
        header = tk.Frame(self, bg=self.BG)
        header.grid(row=0, column=0, sticky="ew", padx=12, pady=(10, 2))
        header.columnconfigure(0, weight=1)
        tk.Label(header, text="Pucoding 用量", bg=self.BG, fg=self.TEXT, font=("Segoe UI", 13, "bold")).grid(row=0, column=0, sticky="w")
        self.refresh_btn = ttk.Button(header, text="↻", width=3, style="Widget.TButton", command=self.refresh)
        self.refresh_btn.grid(row=0, column=1, padx=(6, 0))
        self.pin_btn = ttk.Button(header, text="", width=6, style="Widget.TButton", command=self.toggle_topmost)
        self.pin_btn.grid(row=0, column=2, padx=(4, 0))
        ttk.Button(header, text="⚙", width=3, style="Widget.TButton", command=self._show_settings).grid(row=0, column=3, padx=(4, 0))
        self._update_pin_button()

        self.today_card = tk.Frame(self, bg=self.CARD)
        self.today_card.grid(row=1, column=0, sticky="ew", padx=12, pady=(6, 5))
        self.today_card.columnconfigure(0, weight=1)
        tk.Label(self.today_card, text="今日消费", bg=self.CARD, fg=self.MUTED, font=("Segoe UI", 9)).grid(row=0, column=0, sticky="w", padx=12, pady=(8, 0))
        self.today_value = tk.Label(self.today_card, text="$0.00", bg=self.CARD, fg=self.TEXT, font=("Segoe UI", 24, "bold"))
        self.today_value.grid(row=1, column=0, sticky="w", padx=12, pady=(0, 8))
        self.status_dot = tk.Label(self.today_card, text="●", bg=self.CARD, fg=self.MUTED, font=("Segoe UI", 9))
        self.status_dot.grid(row=0, column=1, padx=(0, 3), pady=(8, 0))
        self.status = tk.Label(self.today_card, text="未连接", bg=self.CARD, fg=self.MUTED, font=("Segoe UI", 8))
        self.status.grid(row=0, column=2, padx=(0, 10), pady=(8, 0))

        details = tk.Frame(self, bg=self.BG)
        details.grid(row=2, column=0, sticky="ew", padx=12, pady=(0, 5))
        details.columnconfigure(0, weight=1)
        details.columnconfigure(1, weight=1)
        self.lifetime_value, self.requests_value = self._metric_card(details, 0, "累计消费", "$0.00")
        self.remaining_value, self.updated_value = self._metric_card(details, 1, "剩余额度", "—")
        self._set_status("等待查询", self.MUTED)
        foot = tk.Frame(self, bg=self.BG)
        foot.grid(row=3, column=0, sticky="ew", padx=12, pady=(0, 8))
        foot.columnconfigure(0, weight=1)
        self.updated_label = tk.Label(foot, text="每 60 秒自动刷新", bg=self.BG, fg=self.MUTED, font=("Segoe UI", 7))
        self.updated_label.grid(row=0, column=0, sticky="w")
        self.requests_label = tk.Label(foot, text="— 次调用", bg=self.BG, fg=self.MUTED, font=("Segoe UI", 7))
        self.requests_label.grid(row=0, column=1, sticky="e")
        self._settings_panel = None

    def _metric_card(self, parent: tk.Widget, column: int, title: str, value: str):
        card = tk.Frame(parent, bg=self.CARD, padx=10, pady=7)
        card.grid(row=0, column=column, sticky="nsew", padx=(0 if column == 0 else 5, 5 if column == 0 else 0))
        tk.Label(card, text=title, bg=self.CARD, fg=self.MUTED, font=("Segoe UI", 8)).pack(anchor="w")
        value_label = tk.Label(card, text=value, bg=self.CARD, fg=self.TEXT, font=("Segoe UI", 13, "bold"))
        value_label.pack(anchor="w", pady=(2, 0))
        sub_label = tk.Label(card, text="", bg=self.CARD, fg=self.MUTED, font=("Segoe UI", 7))
        sub_label.pack(anchor="w")
        return value_label, sub_label

    def _set_status(self, text: str, color: str) -> None:
        self.status.config(text=text, fg=color)
        self.status_dot.config(fg=color)

    def toggle_topmost(self) -> None:
        self._topmost = not self._topmost
        self.attributes("-topmost", self._topmost)
        self._update_pin_button()

    def _update_pin_button(self) -> None:
        if not hasattr(self, "pin_btn"):
            return
        self.pin_btn.config(text="取消置顶" if self._topmost else "置顶")

    def _show_settings(self) -> None:
        if self._settings_panel and self._settings_panel.winfo_exists():
            self._settings_panel.lift()
            return
        panel = tk.Toplevel(self)
        self._settings_panel = panel
        panel.title("Pucoding API Key")
        panel.configure(bg=self.BG)
        panel.resizable(False, False)
        panel.transient(self)
        panel.grab_set()
        tk.Label(panel, text="设置 API Key", bg=self.BG, fg=self.TEXT, font=("Segoe UI", 14, "bold")).pack(anchor="w", padx=18, pady=(18, 3))
        tk.Label(panel, text="仅保存在本机当前用户的 AppData 中", bg=self.BG, fg=self.MUTED, font=("Segoe UI", 9)).pack(anchor="w", padx=18, pady=(0, 12))
        entry = ttk.Entry(panel, width=42, show="•", style="Widget.TEntry")
        entry.insert(0, self.api_key)
        entry.pack(padx=18, fill="x")
        hint = tk.Label(panel, text="Key 应以 sk- 开头；不会出现在请求 URL 中。", bg=self.BG, fg=self.MUTED, font=("Segoe UI", 8))
        hint.pack(anchor="w", padx=18, pady=(7, 14))
        actions = tk.Frame(panel, bg=self.BG)
        actions.pack(fill="x", padx=18, pady=(0, 18))
        ttk.Button(actions, text="取消", style="Widget.TButton", command=panel.destroy).pack(side="right")

        def save() -> None:
            value = entry.get().strip()
            if value and not value.startswith("sk-"):
                hint.config(text="格式看起来不对：Key 应以 sk- 开头。", fg=self.RED)
                return
            self.api_key = value
            save_key(value)
            panel.destroy()
            self._set_status("正在读取…", self.MUTED) if value else self._set_status("请先设置 API Key", self.ORANGE)
            if value:
                self.refresh()

        ttk.Button(actions, text="保存并查询", style="Accent.TButton", command=save).pack(side="right", padx=(0, 8))
        entry.focus_set()

    def refresh(self) -> None:
        if self._loading or not self.api_key:
            return
        if self._auto_refresh_job is not None:
            self.after_cancel(self._auto_refresh_job)
            self._auto_refresh_job = None
        self._loading = True
        self.refresh_btn.config(state="disabled")
        self._set_status("正在读取…", self.MUTED)
        key_snapshot = self.api_key

        def worker() -> None:
            try:
                result = read_usage(key_snapshot)
                self._result_queue.put(("success", result, key_snapshot))
            except Exception as exc:  # noqa: BLE001 - convert all network errors to a short UI message
                self._result_queue.put(("error", str(exc), key_snapshot))

        threading.Thread(target=worker, daemon=True).start()
        if self._poll_job is None:
            self._poll_job = self.after(80, self._poll_result)

    def _poll_result(self) -> None:
        self._poll_job = None
        if self._closing:
            return
        try:
            outcome, payload, key_snapshot = self._result_queue.get_nowait()
        except queue.Empty:
            if self._loading:
                self._poll_job = self.after(80, self._poll_result)
            return

        if key_snapshot == self.api_key:
            if outcome == "success" and isinstance(payload, dict):
                self._apply_usage(payload)
            else:
                self._show_error(str(payload))
        self._finish_refresh()

    def _finish_refresh(self) -> None:
        self._loading = False
        self.refresh_btn.config(state="normal")
        if self.api_key and not self._closing:
            self._auto_refresh_job = self.after(POLL_MS, self.refresh)

    def _apply_usage(self, result: dict) -> None:
        self.today_value.config(text=f"${result['today_cost']:.4f}")
        self.lifetime_value.config(text=f"${result['lifetime_cost']:.4f}")
        self.requests_value.config(text=f"{result['today_requests']:,} 次")
        if result["remaining"] is None or result["quota_limit"] <= 0:
            self.remaining_value.config(text="无限制")
        else:
            self.remaining_value.config(text=f"${result['remaining']:.4f}")
        now = datetime.now().strftime("%H:%M:%S")
        self.updated_label.config(text=f"更新于 {now} · 每 60 秒自动刷新")
        self.requests_label.config(text=f"累计 {result['request_count']:,} 次调用")
        self._set_status(result["status"], self.ACCENT)

    def _show_error(self, message: str) -> None:
        detail = message or "未知错误"
        self._set_status("查询失败", self.RED)
        self.updated_label.config(text=f"{detail[:25]} · 点击 ↻ 重试")

    def _close(self) -> None:
        self._closing = True
        for job in (self._poll_job, self._auto_refresh_job):
            if job is not None:
                try:
                    self.after_cancel(job)
                except tk.TclError:
                    pass
        self.destroy()


def run_self_test() -> int:
    """Verify that the packaged Tcl/Tk runtime can create a window."""
    root = tk.Tk()
    root.withdraw()
    root.update_idletasks()
    root.destroy()
    return 0


if __name__ == "__main__":
    if "--self-test" in sys.argv:
        raise SystemExit(run_self_test())
    app = Widget()
    app.mainloop()