deepseek-tray/main.go

package main

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

	"github.com/getlantern/systray"
)

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

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

// ---------- 配置（优先级：环境变量 > .env 文件 > 默认值） ----------
var (
	apiKey         string
	balanceURL     string
	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() {
	// 加载 .env 文件（同级目录），不覆盖已有环境变量
	loadEnvFile()

	// 读取配置
	apiKey = getEnv("DEEPSEEK_API_KEY", "")
	balanceURL = getEnv("DEEPSEEK_BALANCE_URL", "https://api.deepseek.com/user/balance")

	// 始终启动托盘（即使无 Key，也显示图标+提示）
	systray.Run(onReady, onExit)
}

func onReady() {
	systray.SetIcon(iconData)
	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("低余额告警: 关闭", "切换低余额告警")
	systray.AddSeparator()
	mConfigPath := systray.AddMenuItem("打开 .env 配置文件", "用记事本编辑 .env")
	mAbout := systray.AddMenuItem("关于", "关于本程序")
	mQuit := systray.AddMenuItem("退出", "关闭程序")

	// 首次查询
	go func() {
		if apiKey == "" {
			mBalance.SetTitle("余额: 未配置 Key")
	mBalance.SetTooltip("点击「打开 .env 配置文件」设置 API Key")
		} else {
			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:
					if apiKey != "" {
						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:
			if apiKey == "" {
	showMessageBox("DeepSeek Tray", "未配置 API Key。\n请点击右键菜单「打开 .env 配置文件」设置 DEEPSEEK_API_KEY。")
			} else {
				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))
				if apiKey != "" {
					checkLowBalance()
				}
			} else {
				mThreshold.SetTitle("低余额告警: 关闭")
				mThreshold.SetTooltip("切换低余额告警")
			}

		case <-mConfigPath.ClickedCh:
			openEnvFile()

		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)

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

// ---------- .env 文件支持 ----------

// loadEnvFile 从可执行文件同级目录读取 .env 文件（不覆盖已有环境变量）
func loadEnvFile() {
	exePath, err := os.Executable()
	if err != nil {
		return
	}
	envPath := filepath.Join(filepath.Dir(exePath), ".env")
	data, err := os.ReadFile(envPath)
	if err != nil {
		return // .env 不存在，静默跳过
	}

	lines := strings.Split(string(data), "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			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) == "" {
			os.Setenv(key, val)
		}
	}
}

// openEnvFile 用记事本打开 .env 文件，若不存在则创建模板
func openEnvFile() {
	exePath, err := os.Executable()
	if err != nil {
		showMessageBox("错误", "无法获取程序路径")
		return
	}
	envPath := filepath.Join(filepath.Dir(exePath), ".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())
			return
		}
	}

	// 调用记事本打开
	runNotepad(envPath)
}

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

func getEnv(key, defaultVal string) string {
	// 先查环境变量，再查 .env（已在 loadEnvFile 中注入到 os.Environ）
	if val, ok := os.LookupEnv(key); ok && val != "" {
		return val
	}
	return defaultVal
}

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