deepseek-tray/main.go

package main

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

	"github.com/getlantern/systray"

	_ "embed" // 仅用于 //go:embed 指令
)

//go:embed icon.ico
var appIcon []byte

const (
	httpTimeout     = 15 * time.Second
	refreshMinutes  = 10                         // 自动刷新间隔（分钟）
	lowBalanceAlert = 1.0                        // 低余额告警阈值（元 CNY）
	defaultAPIURL   = "https://api.deepseek.com/user/balance"
)

// main 中写入，之后所有 goroutine 只读。
var (
	apiKey     string
	balanceURL string
	exeDir     string
)

// ---------- API ----------

type balanceAPI struct {
	IsAvailable  bool          `json:"is_available"`
	BalanceInfos []balanceItem `json:"balance_infos"`
}

type balanceItem struct {
	Currency        string `json:"currency"`
	TotalBalance    string `json:"total_balance"`
	GrantedBalance  string `json:"granted_balance"`
	ToppedUpBalance string `json:"topped_up_balance"`
}

// balanceResult 是查询结果的不可变快照。
// ok 为 false 时 text 是错误消息；为 true 时 text 是纯数字金额（不含币种）。
type balanceResult struct {
	ok       bool
	text     string
	currency string
	cny      float64 // CNY 总额，失败时为 0
}

func main() {
	relaunchInUserSession()
	exeDir = executableDir()
	apiKey, balanceURL = loadConfig()
	systray.Run(onReady, func() {})
}

func onReady() {
	systray.SetIcon(appIcon)
	systray.SetTitle("DeepSeek")

	if apiKey != "" {
		systray.SetTooltip("DeepSeek 余额查询 - 正在获取...")
	} else {
		systray.SetTooltip("DeepSeek - 未配置 API Key\n请在 .env 文件中设置 DEEPSEEK_API_KEY")
	}

	mBalance := systray.AddMenuItem("余额: 查询中...", "")
	mBalance.Disable()
	mRefresh := systray.AddMenuItem("刷新余额", "手动刷新")
	mAuto := systray.AddMenuItem("自动刷新: 开启", "切换自动刷新")
	mAlert := systray.AddMenuItem("低余额告警: 关闭", "切换低余额告警")
	systray.AddSeparator()
	mDotEnv := systray.AddMenuItem("打开 .env 配置文件", "编辑 API Key")
	mAbout := systray.AddMenuItem("关于", "状态信息")
	mQuit := systray.AddMenuItem("退出", "关闭程序")

	autoRefresh := true
	alertOn := false

	tickDuration := refreshMinutes * time.Minute

	balanceCh := make(chan balanceResult, 1)
	var ticker *time.Ticker    // 仅 startTicker 写入
	var tickerStop chan struct{}

	startTicker := func() {
		ticker = time.NewTicker(tickDuration)
		tickerStop = make(chan struct{})
		stop := tickerStop

		go func() {
			for {
				select {
				case <-ticker.C:
					if apiKey == "" {
						continue
					}
					r := queryAndUpdate(mBalance)
					select {
					case balanceCh <- r:
					default:
					}
				case <-stop:
					ticker.Stop()
					return
				}
			}
		}()
	}

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

	startTicker()

	// 首次查询（结果同样送入 balanceCh，确保告警逻辑一致）
	go func() {
		if apiKey == "" {
			mBalance.SetTitle("余额: 未配置 Key")
			mBalance.SetTooltip("点击「打开 .env 配置文件」设置 API Key")
			return
		}
		r := queryAndUpdate(mBalance)
		select {
		case balanceCh <- r:
		default:
		}
	}()

	for {
		select {
		case r := <-balanceCh:
			if alertOn && r.cny > 0 && r.cny < lowBalanceAlert {
				alertLowBalance(r.cny)
			}

		case <-mRefresh.ClickedCh:
			if apiKey == "" {
				showMessageBox("DeepSeek Tray",
					"未配置 API Key。\n请点击「打开 .env 配置文件」设置 DEEPSEEK_API_KEY。")
				continue
			}
			queryAndUpdate(mBalance)

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

		case <-mAlert.ClickedCh:
			alertOn = !alertOn
			if alertOn {
				mAlert.SetTitle("低余额告警: 开启")
				if apiKey != "" {
					r := queryAndUpdate(mBalance)
					if r.cny > 0 && r.cny < lowBalanceAlert {
						alertLowBalance(r.cny)
					}
				}
			} else {
				mAlert.SetTitle("低余额告警: 关闭")
			}

		case <-mDotEnv.ClickedCh:
			openDotEnvFile()

		case <-mAbout.ClickedCh:
			showAbout(autoRefresh, alertOn)

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

// ---------- 查询 ----------

func queryAndUpdate(m *systray.MenuItem) balanceResult {
	r := queryBalance()
	now := time.Now().Format("15:04:05")

	if r.ok {
		display := r.text + " " + r.currency
		tip := fmt.Sprintf("余额: %s | %s", display, now)
		systray.SetTooltip(tip)
		m.SetTitle("余额: " + display)
		m.SetTooltip(tip)
		systray.SetTitle(display)
	} else {
		tip := fmt.Sprintf("%s | %s", r.text, now)
		systray.SetTooltip(tip)
		m.SetTitle(r.text)
		m.SetTooltip(tip)
		systray.SetTitle("DeepSeek")
	}
	return r
}

func queryBalance() balanceResult {
	client := &http.Client{Timeout: httpTimeout}
	req, err := http.NewRequest(http.MethodGet, balanceURL, nil)
	if err != nil {
		return balanceResult{text: "创建请求失败"}
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Accept", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return balanceResult{text: "网络错误: " + err.Error()}
	}
	defer resp.Body.Close()

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

	var api balanceAPI
	if err := json.NewDecoder(resp.Body).Decode(&api); err != nil {
		return balanceResult{text: "解析响应失败"}
	}
	if !api.IsAvailable || len(api.BalanceInfos) == 0 {
		return balanceResult{text: "余额不可用"}
	}

	var cnyTotal float64
	for _, item := range api.BalanceInfos {
		if item.Currency == "CNY" {
			var v float64
			fmt.Sscanf(item.TotalBalance, "%f", &v)
			cnyTotal += v
		}
	}

	first := api.BalanceInfos[0]
	return balanceResult{
		ok:       true,
		text:     first.TotalBalance,
		currency: first.Currency,
		cny:      cnyTotal,
	}
}

func alertLowBalance(cny float64) {
	showMessageBox("DeepSeek 余额不足",
		fmt.Sprintf("当前余额 ¥%.2f，低于阈值 ¥%.0f，请及时充值。", cny, lowBalanceAlert))
}

func showAbout(autoRefresh, alertOn bool) {
	showMessageBox("DeepSeek Balance Tray",
		fmt.Sprintf("自动刷新: %s\n刷新间隔: %d 分钟\n低余额告警: %s (阈值 ¥%.0f)",
			yesNo(autoRefresh), refreshMinutes, yesNo(alertOn), lowBalanceAlert))
}

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

// ---------- 配置加载 ----------

// loadConfig 返回 (apiKey, balanceURL)。
// apiKey 优先级：DEEPSEEK_API_KEY 环境变量 > .env 文件。
// balanceURL 优先级：DEEPSEEK_BALANCE_URL 环境变量 > .env 文件 > defaultAPIURL。
func loadConfig() (key, url string) {
	dotenv := readDotEnv(filepath.Join(exeDir, ".env"))
	// dotenv 可能为 nil；Go 读取 nil map 返回零值，安全。

	if v := os.Getenv("DEEPSEEK_API_KEY"); v != "" {
		key = v
	} else if dotenv != nil {
		key = dotenv["DEEPSEEK_API_KEY"]
	}

	if v := os.Getenv("DEEPSEEK_BALANCE_URL"); v != "" {
		url = v
	} else if dotenv != nil {
		if v := dotenv["DEEPSEEK_BALANCE_URL"]; v != "" {
			url = v
		}
	}
	if url == "" {
		url = defaultAPIURL
	}
	return
}

func readDotEnv(path string) map[string]string {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil
	}
	m := make(map[string]string)
	for _, raw := range strings.Split(string(data), "\n") {
		line := strings.TrimSpace(raw)
		if line == "" || line[0] == '#' {
			continue
		}
		eq := strings.IndexByte(line, '=')
		if eq < 1 {
			continue
		}
		k := strings.TrimSpace(line[:eq])
		v := strings.TrimSpace(line[eq+1:])
		v = strings.Trim(v, `"'`)
		m[k] = v
	}
	return m
}

func openDotEnvFile() {
	path := filepath.Join(exeDir, ".env")

	if _, err := os.Stat(path); os.IsNotExist(err) {
		const 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(path)
}

// ---------- 工具 ----------

func executableDir() string {
	exe, err := os.Executable()
	if err != nil {
		wd, _ := os.Getwd()
		return wd
	}
	return filepath.Dir(exe)
}