deepseek-tray/main.go

package main

import (
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/getlantern/systray"
)

// ---------- 嵌入资源 ----------
//go:embed icon.ico
var iconData []byte

// 确保 embed 包被"使用"（Go 1.26+ 要求）
var _ = embed.FS{}

// ---------- 默认配置（可被环境变量覆盖） ----------
var (
	apiKey         = getEnv("DEEPSEEK_API_KEY", "")
	balanceURL     = getEnv("DEEPSEEK_BALANCE_URL", "https://api.deepseek.com/user/balance")
	requestTimeout = 15 // 秒
	refreshMinutes = 10 // 自动刷新间隔，0 表示关闭
	lowThreshold   = 1.0 // 余额低于此值（元）时告警
)

// ---------- 运行时状态 ----------
var (
	lastBalance string
	lastUpdate  time.Time
	mu          sync.RWMutex
)

// ---------- 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"`
}

func main() {
	if apiKey == "" {
		fmt.Fprintln(os.Stderr, "错误: 环境变量 DEEPSEEK_API_KEY 未设置")
		fmt.Fprintln(os.Stderr, "请设置后运行: set DEEPSEEK_API_KEY=sk-your-key")
		os.Exit(1)
	}
	systray.Run(onReady, onExit)
}

func onReady() {
	systray.SetIcon(iconData)
	systray.SetTitle("DeepSeek")
	systray.SetTooltip("DeepSeek 余额查询 - 正在获取...")

	// 菜单项
	mBalance := systray.AddMenuItem("余额: 查询中...", "")
	mBalance.Disable()

	mManualRefresh := systray.AddMenuItem("刷新余额", "手动刷新")
	mAutoRefresh := systray.AddMenuItem("自动刷新: 开启", "切换自动刷新")
	mThreshold := systray.AddMenuItem("低余额告警: 关闭", "切换低余额告警")
	systray.AddSeparator()
	mAbout := systray.AddMenuItem("关于", "关于本程序")
	mQuit := systray.AddMenuItem("退出", "关闭程序")

	// 首次查询
	go refreshBalance(mBalance)

	// 定时器控制
	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{})
		go func() {
			for {
				select {
				case <-ticker.C:
					refreshBalance(mBalance)
					if alertOn {
						checkLowBalance()
					}
				case <-tickerStop:
					ticker.Stop()
					return
				}
			}
		}()
	}

	stopTicker := func() {
		if tickerStop != nil {
			close(tickerStop)
			tickerStop = nil
		}
	}

	startTicker()

	// 事件循环
	for {
		select {
		case <-mManualRefresh.ClickedCh:
			refreshBalance(mBalance)

		case <-mAutoRefresh.ClickedCh:
			if autoRefresh {
				stopTicker()
				mAutoRefresh.SetTitle("自动刷新: 关闭")
				autoRefresh = false
			} else {
				startTicker()
				mAutoRefresh.SetTitle("自动刷新: 开启")
				autoRefresh = true
			}

		case <-mThreshold.ClickedCh:
			alertOn = !alertOn
			if alertOn {
				mThreshold.SetTitle("低余额告警: 开启")
				mThreshold.SetTooltip(fmt.Sprintf("余额低于 ¥%.0f 时提醒", lowThreshold))
				checkLowBalance()
			} else {
				mThreshold.SetTitle("低余额告警: 关闭")
				mThreshold.SetTooltip("切换低余额告警")
			}

		case <-mAbout.ClickedCh:
			systray.SetTooltip(fmt.Sprintf(
				"自动刷新:%s | 间隔:%dmin | 告警:%s",
				boolToStr(autoRefresh), refreshMinutes, boolToStr(alertOn),
			))

		case <-mQuit.ClickedCh:
			stopTicker()
			systray.Quit()
			return
		}
	}
}

func onExit() {
	// 清理
}

// refreshBalance 请求 API 并更新托盘状态
func refreshBalance(menuItem *systray.MenuItem) {
	text, currency, totalVal := fetchBalance()
	mu.Lock()
	lastBalance = text
	lastUpdate = time.Now()
	mu.Unlock()

	tooltip := fmt.Sprintf("DeepSeek 余额: %s %s\n更新时间: %s",
		text, currency, lastUpdate.Format("15:04:05"))
	systray.SetTooltip(tooltip)
	menuItem.SetTitle("余额: " + text + " " + currency)
	menuItem.SetTooltip(tooltip)

	// 更新程序标题
	if totalVal > 0 {
		systray.SetTitle(text)
	} else {
		systray.SetTitle("DeepSeek")
	}
}

// fetchBalance 返回 (余额文本, 货币, 数值用于比较)
func fetchBalance() (string, string, float64) {
	client := &http.Client{Timeout: time.Duration(requestTimeout) * time.Second}
	req, err := http.NewRequest("GET", balanceURL, nil)
	if err != nil {
		return "创建请求失败", "", 0
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Accept", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return "请求失败: " + err.Error(), "", 0
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusUnauthorized {
		return "API Key 无效", "", 0
	}
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(body)), "", 0
	}

	var br balanceResponse
	if err := json.NewDecoder(resp.Body).Decode(&br); err != nil {
		return "解析响应失败", "", 0
	}

	if !br.IsAvailable || len(br.BalanceInfos) == 0 {
		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
	}

	if len(parts) == 1 {
		return parts[0], currency, total
	}
	return fmt.Sprintf("%v", parts), currency, total
}

// checkLowBalance 检查余额是否低于阈值，通过弹窗告警（使用 Windows API）
func checkLowBalance() {
	_, _, totalVal := fetchBalance()
	if totalVal > 0 && totalVal < lowThreshold {
		msg := fmt.Sprintf("当前余额 ¥%.2f，低于阈值 ¥%.0f，请及时充值。", totalVal, lowThreshold)
		showMessageBox("⚠️ DeepSeek 余额不足", msg)
	}
}

// ---------- 工具函数 ----------

func getEnv(key, defaultVal string) string {
	if val, ok := os.LookupEnv(key); ok && val != "" {
		return val
	}
	return defaultVal
}

func boolToStr(b bool) string {
	if b {
		return "是"
	}
	return "否"
}