Explorar el Código

update main.go

caesar hace 2 meses
padre
commit
d1163404ce
Se han modificado 1 ficheros con 153 adiciones y 181 borrados
  1. 153 181
      main.go

+ 153 - 181
main.go

@@ -9,115 +9,93 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/getlantern/systray"
 )
 
-// ---------- 嵌入资源 ----------
 //go:embed icon.ico
-var iconData []byte
+var icon []byte
 
-// 确保 embed 包被"使用"(Go 1.26+ 要求)
-var _ = embed.FS{}
+var _ = embed.FS{} // 保证 Go 1.26+ embed 引用检查
 
-// ---------- 配置(优先级:环境变量 > .env 文件 > 默认值) ----------
-var (
-	apiKey         string
-	balanceURL     string
-	requestTimeout = 15 // 秒
-	refreshMinutes = 10 // 自动刷新间隔,0 表示关闭
-	lowThreshold   = 1.0 // 余额低于此值(元)时告警
+// ---------- 配置常量 ----------
+const (
+	apiTimeout       = 15 * time.Second  // HTTP 请求超时
+	refreshPeriod    = 10 * time.Minute   // 自动刷新间隔
+	lowBalanceYuan   = 1.0                // 低余额告警阈值(元)
+	balanceDefaultURL = "https://api.deepseek.com/user/balance"
 )
 
-// ---------- 运行时状态 ----------
 var (
-	lastBalance string
-	lastUpdate  time.Time
-	mu          sync.RWMutex
+	apiKey     string
+	balanceURL = balanceDefaultURL
 )
 
-// ---------- API 响应结构 ----------
-type balanceResponse struct {
-	IsAvailable  bool `json:"is_available"`
-	BalanceInfos []struct {
-		Currency        string `json:"currency"`
-		TotalBalance    string `json:"total_balance"`
-		GrantedBalance  string `json:"granted_balance"`
-		ToppedUpBalance string `json:"topped_up_balance"`
-	} `json:"balance_infos"`
+// ---------- API 响应 ----------
+type balanceResp struct {
+	IsAvailable  bool          `json:"is_available"`
+	BalanceInfos []balanceItem `json:"balance_infos"`
 }
 
-func main() {
-	// 加载 .env 文件(同级目录),不覆盖已有环境变量
-	loadEnvFile()
-
-	// 读取配置
-	apiKey = getEnv("DEEPSEEK_API_KEY", "")
-	balanceURL = getEnv("DEEPSEEK_BALANCE_URL", "https://api.deepseek.com/user/balance")
+type balanceItem struct {
+	Currency        string `json:"currency"`
+	TotalBalance    string `json:"total_balance"`
+	GrantedBalance  string `json:"granted_balance"`
+	ToppedUpBalance string `json:"topped_up_balance"`
+}
 
-	// 始终启动托盘(即使无 Key,也显示图标+提示)
-	systray.Run(onReady, onExit)
+// ---------- 入口 ----------
+func main() {
+	loadDotEnv()
+	apiKey = envOr("DEEPSEEK_API_KEY", "")
+	if v := envOr("DEEPSEEK_BALANCE_URL", ""); v != "" {
+		balanceURL = v
+	}
+	systray.Run(onReady, func() {})
 }
 
 func onReady() {
-	systray.SetIcon(iconData)
+	systray.SetIcon(icon)
 	systray.SetTitle("DeepSeek")
-
 	if apiKey == "" {
 		systray.SetTooltip("DeepSeek - 未配置 API Key\n请在 .env 文件中设置 DEEPSEEK_API_KEY")
 	} else {
 		systray.SetTooltip("DeepSeek 余额查询 - 正在获取...")
 	}
 
-	// 菜单
+	// 菜单
 	mBalance := systray.AddMenuItem("余额: 查询中...", "")
 	mBalance.Disable()
-
-	mManualRefresh := systray.AddMenuItem("刷新余额", "手动刷新")
-	mAutoRefresh := systray.AddMenuItem("自动刷新: 开启", "切换自动刷新")
-	mThreshold := systray.AddMenuItem("低余额告警: 关闭", "切换低余额告警")
+	mRefresh := systray.AddMenuItem("刷新余额", "手动刷新")
+	mAuto := systray.AddMenuItem("自动刷新: 开启", "切换自动刷新")
+	mAlert := systray.AddMenuItem("低余额告警: 关闭", "切换低余额告警")
 	systray.AddSeparator()
-	mConfigPath := systray.AddMenuItem("打开 .env 配置文件", "用记事本编辑 .env")
-	mAbout := systray.AddMenuItem("关于", "关于本程序")
+	mDotEnv := systray.AddMenuItem("打开 .env 配置文件", "编辑 API Key")
+	mAbout := systray.AddMenuItem("关于", "状态信息")
 	mQuit := systray.AddMenuItem("退出", "关闭程序")
 
-	// 首次查询
-	go func() {
-		if apiKey == "" {
-			mBalance.SetTitle("余额: 未配置 Key")
-	mBalance.SetTooltip("点击「打开 .env 配置文件」设置 API Key")
-		} else {
-			refreshBalance(mBalance)
-		}
-	}()
+	// 运行时状态(仅在事件循环 goroutine 中访问,无数据竞争)
+	autoRefresh := true
+	alertOn := false
+	var ticker *time.Ticker
+	var tickStop chan struct{}
 
-	// 定时器控制
-	var (
-		ticker      *time.Ticker
-		tickerStop  chan struct{}
-		autoRefresh = true
-		alertOn     = false
-	)
-
-	startTicker := func() {
-		if refreshMinutes <= 0 {
-			return
-		}
-		ticker = time.NewTicker(time.Duration(refreshMinutes) * time.Minute)
-		tickerStop = make(chan struct{})
+	startTick := func() {
+		ticker = time.NewTicker(refreshPeriod)
+		tickStop = make(chan struct{})
 		go func() {
 			for {
 				select {
 				case <-ticker.C:
 					if apiKey != "" {
-						refreshBalance(mBalance)
-						if alertOn {
-							checkLowBalance()
+						balance, _ := refreshAndUpdate(mBalance)
+						if alertOn && balance < lowBalanceYuan && balance > 0 {
+	showMessageBox("DeepSeek 余额不足",
+								fmt.Sprintf("当前余额 ¥%.2f,低于阈值 ¥%.0f,请及时充值。", balance, lowBalanceYuan))
 						}
 					}
-				case <-tickerStop:
+				case <-tickStop:
 					ticker.Stop()
 					return
 				}
@@ -125,94 +103,110 @@ func onReady() {
 		}()
 	}
 
-	stopTicker := func() {
-		if tickerStop != nil {
-			close(tickerStop)
-			tickerStop = nil
+	stopTick := func() {
+		if tickStop != nil {
+			close(tickStop)
+			tickStop = nil
 		}
 	}
 
-	startTicker()
+	startTick()
+
+	// 首次查询
+	go func() {
+		if apiKey == "" {
+			mBalance.SetTitle("余额: 未配置 Key")
+			mBalance.SetTooltip("点击「打开 .env 配置文件」设置 API Key")
+		} else {
+			refreshAndUpdate(mBalance)
+		}
+	}()
 
 	// 事件循环
 	for {
 		select {
-		case <-mManualRefresh.ClickedCh:
+		case <-mRefresh.ClickedCh:
 			if apiKey == "" {
-	showMessageBox("DeepSeek Tray", "未配置 API Key。\n请点击右键菜单「打开 .env 配置文件」设置 DEEPSEEK_API_KEY。")
+	showMessageBox("DeepSeek Tray",
+					"未配置 API Key。\n请点击「打开 .env 配置文件」设置 DEEPSEEK_API_KEY。")
 			} else {
-				refreshBalance(mBalance)
+				refreshAndUpdate(mBalance)
 			}
 
-		case <-mAutoRefresh.ClickedCh:
+		case <-mAuto.ClickedCh:
 			if autoRefresh {
-				stopTicker()
-				mAutoRefresh.SetTitle("自动刷新: 关闭")
-				autoRefresh = false
+				stopTick()
+				mAuto.SetTitle("自动刷新: 关闭")
 			} else {
-				startTicker()
-				mAutoRefresh.SetTitle("自动刷新: 开启")
-				autoRefresh = true
+				startTick()
+				mAuto.SetTitle("自动刷新: 开启")
 			}
+			autoRefresh = !autoRefresh
 
-		case <-mThreshold.ClickedCh:
+		case <-mAlert.ClickedCh:
 			alertOn = !alertOn
 			if alertOn {
-				mThreshold.SetTitle("低余额告警: 开启")
-				mThreshold.SetTooltip(fmt.Sprintf("余额低于 ¥%.0f 时提醒", lowThreshold))
+				mAlert.SetTitle("低余额告警: 开启")
+				mAlert.SetTooltip(fmt.Sprintf("余额低于 ¥%.0f 时提醒", lowBalanceYuan))
 				if apiKey != "" {
-					checkLowBalance()
+					balance, _ := refreshAndUpdate(mBalance)
+					if balance < lowBalanceYuan && balance > 0 {
+						showMessageBox("DeepSeek 余额不足",
+							fmt.Sprintf("当前余额 ¥%.2f,低于阈值 ¥%.0f,请及时充值。", balance, lowBalanceYuan))
+					}
 				}
 			} else {
-				mThreshold.SetTitle("低余额告警: 关闭")
-				mThreshold.SetTooltip("切换低余额告警")
+				mAlert.SetTitle("低余额告警: 关闭")
+				mAlert.SetTooltip("切换低余额告警")
 			}
 
-		case <-mConfigPath.ClickedCh:
-			openEnvFile()
+		case <-mDotEnv.ClickedCh:
+			openDotEnv()
 
 		case <-mAbout.ClickedCh:
-			info := fmt.Sprintf("DeepSeek 余额查询系统托盘\n\n自动刷新: %s\n刷新间隔: %d 分钟\n低余额告警: %s\n\n配置文件: deepseek-tray.exe 同目录下的 .env",
-				boolToStr(autoRefresh), refreshMinutes, boolToStr(alertOn))
-			showMessageBox("DeepSeek Balance Tray", info)
+			ar := "否"
+			if autoRefresh {
+				ar = "是"
+			}
+			ao := "否"
+			if alertOn {
+				ao = "是"
+			}
+	showMessageBox("DeepSeek Balance Tray",
+				fmt.Sprintf("DeepSeek 余额查询系统托盘\n\n自动刷新: %s\n刷新间隔: %d 分钟\n低余额告警: %s\n\n配置文件: .env",
+					ar, int(refreshPeriod.Minutes()), ao))
 
 		case <-mQuit.ClickedCh:
-			stopTicker()
+			stopTick()
 			systray.Quit()
 			return
 		}
 	}
 }
 
-func onExit() {
-	// 清理
-}
+// ---------- 核心:余额获取与 UI 更新 ----------
 
-// refreshBalance 请求 API 并更新托盘状态
-func refreshBalance(menuItem *systray.MenuItem) {
-	text, currency, totalVal := fetchBalance()
-	mu.Lock()
-	lastBalance = text
-	lastUpdate = time.Now()
-	mu.Unlock()
+// refreshAndUpdate 请求 API 并更新托盘状态,返回 CNY 余额数值。
+func refreshAndUpdate(m *systray.MenuItem) (cnyBalance float64, err error) {
+	text, currency, cny := fetchBalance()
 
 	tooltip := fmt.Sprintf("DeepSeek 余额: %s %s\n更新时间: %s",
-		text, currency, lastUpdate.Format("15:04:05"))
+		text, currency, time.Now().Format("15:04:05"))
 	systray.SetTooltip(tooltip)
-	menuItem.SetTitle("余额: " + text + " " + currency)
-	menuItem.SetTooltip(tooltip)
+	m.SetTitle("余额: " + text + " " + currency)
+	m.SetTooltip(tooltip)
 
-	// 更新程序标题(显示在托盘旁)
-	if totalVal > 0 {
+	if cny > 0 {
 		systray.SetTitle(text)
 	} else {
 		systray.SetTitle("DeepSeek")
 	}
+	return cny, nil
 }
 
-// fetchBalance 返回 (余额文本, 货币, 数值用于比较)
+// fetchBalance 返回 (余额文本, 货币, CNY 数值)
 func fetchBalance() (string, string, float64) {
-	client := &http.Client{Timeout: time.Duration(requestTimeout) * time.Second}
+	client := &http.Client{Timeout: apiTimeout}
 	req, err := http.NewRequest("GET", balanceURL, nil)
 	if err != nil {
 		return "创建请求失败", "", 0
@@ -234,7 +228,7 @@ func fetchBalance() (string, string, float64) {
 		return fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(body)), "", 0
 	}
 
-	var br balanceResponse
+	var br balanceResp
 	if err := json.NewDecoder(resp.Body).Decode(&br); err != nil {
 		return "解析响应失败", "", 0
 	}
@@ -243,106 +237,84 @@ func fetchBalance() (string, string, float64) {
 		return "余额不可用", "", 0
 	}
 
-	// 合并所有币种
-	var parts []string
-	var total float64
-	currency := "CNY"
-	for _, info := range br.BalanceInfos {
-		var val float64
-		fmt.Sscanf(info.TotalBalance, "%f", &val)
-		parts = append(parts, fmt.Sprintf("%s %s", info.TotalBalance, info.Currency))
-		if info.Currency == "CNY" {
-			total += val
-		}
-		currency = info.Currency
-	}
+	// 取第一个币种为主显示,叠加所有 CNY 为数值
+	info := br.BalanceInfos[0]
+	first := fmt.Sprintf("%s %s", info.TotalBalance, info.Currency)
 
-	if len(parts) == 1 {
-		return parts[0], currency, total
+	var cnyTotal float64
+	for _, item := range br.BalanceInfos {
+		var v float64
+		fmt.Sscanf(item.TotalBalance, "%f", &v)
+		if item.Currency == "CNY" {
+			cnyTotal += v
+		}
 	}
-	return fmt.Sprintf("%v", parts), currency, total
-}
 
-// checkLowBalance 检查余额是否低于阈值,通过弹窗告警
-func checkLowBalance() {
-	_, _, totalVal := fetchBalance()
-	if totalVal > 0 && totalVal < lowThreshold {
-		msg := fmt.Sprintf("当前余额 ¥%.2f,低于阈值 ¥%.0f,请及时充值。", totalVal, lowThreshold)
-		showMessageBox("⚠️ DeepSeek 余额不足", msg)
-	}
+	return first, info.Currency, cnyTotal
 }
 
-// ---------- .env 文件支持 ----------
+// ---------- .env 支持 ----------
 
-// loadEnvFile 从可执行文件同级目录读取 .env 文件(不覆盖已有环境变量)
-func loadEnvFile() {
-	exePath, err := os.Executable()
-	if err != nil {
+func loadDotEnv() {
+	dir := exeDir()
+	if dir == "" {
 		return
 	}
-	envPath := filepath.Join(filepath.Dir(exePath), ".env")
-	data, err := os.ReadFile(envPath)
+	data, err := os.ReadFile(filepath.Join(dir, ".env"))
 	if err != nil {
-		return // .env 不存在,静默跳过
+		return
 	}
-
-	lines := strings.Split(string(data), "\n")
-	for _, line := range lines {
+	for _, line := range strings.Split(string(data), "\n") {
 		line = strings.TrimSpace(line)
-		if line == "" || strings.HasPrefix(line, "#") {
+		if line == "" || line[0] == '#' {
 			continue
 		}
-		// 支持 KEY=VALUE 或 KEY="VALUE" 格式
 		idx := strings.IndexByte(line, '=')
 		if idx < 0 {
 			continue
 		}
 		key := strings.TrimSpace(line[:idx])
-		val := strings.TrimSpace(line[idx+1:])
-		val = strings.Trim(val, `"`) // 移除双引号
-
-		// 不覆盖已有环境变量
-		if os.Getenv(key) == "" {
+		val := strings.Trim(strings.TrimSpace(line[idx+1:]), `"`)
+		// .env 不覆盖已有环境变量
+		if os.Getenv(key) == "" && key != "" {
 			os.Setenv(key, val)
 		}
 	}
 }
 
-// openEnvFile 用记事本打开 .env 文件,若不存在则创建模板
-func openEnvFile() {
-	exePath, err := os.Executable()
-	if err != nil {
-		showMessageBox("错误", "无法获取程序路径")
+func openDotEnv() {
+	dir := exeDir()
+	if dir == "" {
+	showMessageBox("错误", "无法获取程序路径")
 		return
 	}
-	envPath := filepath.Join(filepath.Dir(exePath), ".env")
+	path := filepath.Join(dir, ".env")
 
-	// 如果 .env 不存在,创建模板
-	if _, err := os.Stat(envPath); os.IsNotExist(err) {
-		tmpl := "# DeepSeek API Key(必填)\nDEEPSEEK_API_KEY=sk-your-key-here\n\n# DeepSeek 余额查询接口(可选)\n# DEEPSEEK_BALANCE_URL=https://api.deepseek.com/user/balance\n"
-		if err := os.WriteFile(envPath, []byte(tmpl), 0644); err != nil {
-			showMessageBox("错误", "无法创建 .env 文件: "+err.Error())
+	if _, err := os.Stat(path); os.IsNotExist(err) {
+		tmpl := "# DeepSeek API Key(必填)\nDEEPSEEK_API_KEY=sk-your-key-here\n\n# 余额查询接口(可选)\n# DEEPSEEK_BALANCE_URL=https://api.deepseek.com/user/balance\n"
+		if err := os.WriteFile(path, []byte(tmpl), 0644); err != nil {
+	showMessageBox("错误", "无法创建 .env 文件: "+err.Error())
 			return
 		}
 	}
 
-	// 调用记事本打开
-	runNotepad(envPath)
+	runNotepad(path)
 }
 
-// ---------- 工具函数 ----------
-
-func getEnv(key, defaultVal string) string {
-	// 先查环境变量,再查 .env(已在 loadEnvFile 中注入到 os.Environ)
-	if val, ok := os.LookupEnv(key); ok && val != "" {
-		return val
+// exeDir 返回可执行文件所在目录路径。
+func exeDir() string {
+	p, err := os.Executable()
+	if err != nil {
+		return ""
 	}
-	return defaultVal
+	return filepath.Dir(p)
 }
 
-func boolToStr(b bool) string {
-	if b {
-		return "是"
+// ---------- 工具 ----------
+
+func envOr(key, def string) string {
+	if v, ok := os.LookupEnv(key); ok && v != "" {
+		return v
 	}
-	return "否"
+	return def
 }