deepseek-tray/balance.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"
)

const (
	balanceHTTPTimeout  = 15 * time.Second
	balanceRefreshMin   = 10  // auto-refresh interval in minutes
	lowBalanceThreshold = 1.0 // alert when CNY total drops below this (元)
)

// envTemplate is written to .env when the user opens a missing config file.
const envTemplate = "# DeepSeek API Key（必填）\nDEEPSEEK_API_KEY=sk-your-key-here\n\n# 余额查询接口（可选）\n# DEEPSEEK_BALANCE_URL=https://api.deepseek.com/user/balance\n"

// BalanceAPI abstracts the balance endpoint so the tray feature is testable
// without the network and without globals.
type BalanceAPI interface {
	Query(ctx context.Context, apiKey string) (Balance, error)
}

// Balance is a successful balance response.
type Balance struct {
	Available bool
	Items     []BalanceItem
}

// BalanceItem is one currency row from the DeepSeek balance endpoint.
type BalanceItem struct {
	Currency        string `json:"currency"`
	TotalBalance    string `json:"total_balance"`
	GrantedBalance  string `json:"granted_balance"`
	ToppedUpBalance string `json:"topped_up_balance"`
}

type balanceResponse struct {
	IsAvailable  bool          `json:"is_available"`
	BalanceInfos []BalanceItem `json:"balance_infos"`
}

// httpBalanceAPI is the production BalanceAPI implementation.
type httpBalanceAPI struct {
	URL    string
	Client *http.Client
}

func (a httpBalanceAPI) Query(ctx context.Context, apiKey string) (Balance, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, a.URL, nil)
	if err != nil {
		return Balance{}, fmt.Errorf("创建请求失败: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Accept", "application/json")

	resp, err := a.Client.Do(req)
	if err != nil {
		return Balance{}, fmt.Errorf("网络错误: %w", err)
	}
	defer resp.Body.Close()

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

	var api balanceResponse
	if err := json.NewDecoder(resp.Body).Decode(&api); err != nil {
		return Balance{}, fmt.Errorf("解析响应失败: %w", err)
	}
	if !api.IsAvailable || len(api.BalanceInfos) == 0 {
		return Balance{}, fmt.Errorf("余额不可用")
	}
	return Balance{Available: true, Items: api.BalanceInfos}, nil
}

// Display returns the primary display amount and its currency.
func (b Balance) Display() (text, currency string) {
	if len(b.Items) == 0 {
		return "", ""
	}
	return b.Items[0].TotalBalance, b.Items[0].Currency
}

// CNYTotal sums every CNY row, used by the low-balance alert.
func (b Balance) CNYTotal() float64 {
	var sum float64
	for _, it := range b.Items {
		if it.Currency == "CNY" {
			var v float64
			fmt.Sscanf(it.TotalBalance, "%f", &v)
			sum += v
		}
	}
	return sum
}

// balanceFeature implements Feature for the DeepSeek balance tray UI.
type balanceFeature struct {
	cfg        Config
	api        BalanceAPI
	dotenvPath string

	autoRefresh bool // owned by the Run goroutine only
	alertOn     bool // owned by the Run goroutine only

	menu *balanceMenu
}

type balanceMenu struct {
	balance *Item
	refresh *Item
	auto    *Item
	alert   *Item
	dotenv  *Item
	about   *Item
	quit    *Item
}

func newBalanceFeature(cfg Config, dotenvPath string) *balanceFeature {
	return &balanceFeature{
		cfg:         cfg,
		api:         httpBalanceAPI{URL: cfg.BalanceURL, Client: &http.Client{Timeout: balanceHTTPTimeout}},
		dotenvPath:  dotenvPath,
		autoRefresh: true,
	}
}

func (b *balanceFeature) Name() string { return "balance" }

func (b *balanceFeature) Install(m *Menu) {
	b.menu = &balanceMenu{
		balance: m.AddItem("余额: 查询中...", ""),
		refresh: m.AddItem("刷新余额", "手动刷新"),
		auto:    m.AddItem("自动刷新: 开启", "切换自动刷新"),
		alert:   m.AddItem("低余额告警: 关闭", "切换低余额告警"),
	}
	b.menu.balance.Disable()
	m.AddSeparator()
	b.menu.dotenv = m.AddItem("打开 .env 配置文件", "编辑 API Key")
	b.menu.about = m.AddItem("关于", "状态信息")
	b.menu.quit = m.AddItem("退出", "关闭程序")

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

func (b *balanceFeature) Stop() {}

func (b *balanceFeature) Run(ctx context.Context) {
	results := make(chan queryResult, 1)

	go func() {
		if b.cfg.APIKey == "" {
			b.menu.balance.SetTitle("余额: 未配置 Key")
			b.menu.balance.SetTooltip("点击「打开 .env 配置文件」设置 API Key")
			return
		}
		b.queryInto(b.menu.balance, results)
	}()

	ticker := time.NewTicker(balanceRefreshMin * time.Minute)
	defer ticker.Stop()

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

		case <-ticker.C:
			if b.cfg.APIKey == "" || !b.autoRefresh {
				continue
			}
			b.queryInto(b.menu.balance, results)

		case <-b.menu.refresh.Clicked:
			if b.cfg.APIKey == "" {
				showMessageBox("DeepSeek Tray", "未配置 API Key。\n请点击「打开 .env 配置文件」设置 DEEPSEEK_API_KEY。")
				continue
			}
			b.queryInto(b.menu.balance, results)

		case <-b.menu.auto.Clicked:
			b.autoRefresh = !b.autoRefresh
			if b.autoRefresh {
				b.menu.auto.SetTitle("自动刷新: 开启")
			} else {
				b.menu.auto.SetTitle("自动刷新: 关闭")
			}

		case <-b.menu.alert.Clicked:
			b.alertOn = !b.alertOn
			if b.alertOn {
				b.menu.alert.SetTitle("低余额告警: 开启")
				if b.cfg.APIKey != "" {
					b.queryInto(b.menu.balance, results)
				}
			} else {
				b.menu.alert.SetTitle("低余额告警: 关闭")
			}

		case <-b.menu.dotenv.Clicked:
			openDotEnvFile(b.dotenvPath)

		case <-b.menu.about.Clicked:
			showAbout(b.autoRefresh, b.alertOn)

		case <-b.menu.quit.Clicked:
			quitTray()
			return

		case <-ctx.Done():
			return
		}
	}
}

// queryResult is a snapshot used for the low-balance alert decision.
type queryResult struct {
	ok  bool
	cny float64
}

// queryInto queries the balance endpoint, updates the tray UI, and reports the
// result for the alert path (non-blocking).
func (b *balanceFeature) queryInto(it *Item, results chan<- queryResult) {
	bal, err := b.api.Query(context.Background(), b.cfg.APIKey)
	now := time.Now().Format("15:04:05")

	var r queryResult
	if err == nil {
		text, currency := bal.Display()
		display := text + " " + currency
		tip := fmt.Sprintf("余额: %s | %s", display, now)
		setTrayTooltip(tip)
		it.SetTitle("余额: " + display)
		it.SetTooltip(tip)
		setTrayTitle(display)
		r = queryResult{ok: true, cny: bal.CNYTotal()}
	} else {
		tip := fmt.Sprintf("%s | %s", err.Error(), now)
		setTrayTooltip(tip)
		it.SetTitle(err.Error())
		it.SetTooltip(tip)
		setTrayTitle("DeepSeek")
	}

	select {
	case results <- r:
	default:
	}
}

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

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

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

// openDotEnvFile opens the .env file, creating a template first when absent.
func openDotEnvFile(path string) {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		if err := os.WriteFile(path, []byte(envTemplate), 0644); err != nil {
			showMessageBox("错误", "无法创建 .env 文件: "+err.Error())
			return
		}
	}
	runNotepad(path)
}