sub2api/backend/internal/service/antigravity_credits.go

package service

import (
	"context"
	"log"
	"strconv"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
)

// AI Credits（GOOGLE_ONE_AI）状态管理：
//   - free tier 耗尽时，可注入 v1internal.enabledCreditTypes=["GOOGLE_ONE_AI"] 落到付费余额。
//   - 账号余额持久化在 Account.Extra，避免每次请求都去 loadCodeAssist 查询。
//   - INSUFFICIENT_G1_CREDITS_BALANCE 错误回写 exhausted_at 时间戳，让请求转换器跳过该账号的 credits 注入。
//   - loadCodeAssist 时主动刷新余额并清除 exhausted 标记。

const (
	// extraKeyCreditsBalance 缓存的 GOOGLE_ONE_AI 可用余额（float64，单位由上游决定）。
	extraKeyCreditsBalance = "antigravity_credits_balance"
	// extraKeyCreditsCheckedAt 余额最近查询时间（RFC3339）。
	extraKeyCreditsCheckedAt = "antigravity_credits_checked_at"
	// extraKeyCreditsExhaustedAt 上次收到 INSUFFICIENT_G1_CREDITS_BALANCE 的时间（RFC3339）。
	extraKeyCreditsExhaustedAt = "antigravity_credits_exhausted_at"

	// creditsExhaustedRecheckInterval 余额耗尽后的重新探测间隔。
	// 在此间隔内不再注入 enabledCreditTypes，避免反复触发 INSUFFICIENT_G1_CREDITS_BALANCE。
	// 间隔到达后允许下一次 loadCodeAssist 刷新余额并解除标记。
	creditsExhaustedRecheckInterval = 30 * time.Minute
)

// AccountHasUsableCredits 判断账号当前是否可注入 enabledCreditTypes。
//   - 仅当最近一次余额查询 > 0 且 exhausted_at 已过期才返回 true。
//   - 从未查询过余额时返回 false（保守策略：不知道有就不注入，避免无效请求）。
func AccountHasUsableCredits(account *Account) bool {
	if account == nil || account.Extra == nil {
		return false
	}

	// 余额耗尽标记仍在生效期内 → 不可用
	if exhaustedAtStr, ok := account.Extra[extraKeyCreditsExhaustedAt].(string); ok && exhaustedAtStr != "" {
		if t, err := time.Parse(time.RFC3339, exhaustedAtStr); err == nil {
			if time.Since(t) < creditsExhaustedRecheckInterval {
				return false
			}
		}
	}

	balance := readFloat(account.Extra[extraKeyCreditsBalance])
	return balance > 0
}

// refreshAccountCreditsFromLoadCodeAssist 从 loadCodeAssist 响应里提取 paidTier.availableCredits。
// 任何 GOOGLE_ONE_AI 类型的余额（或 creditType 为空时按总额）都会被写入 Account.Extra。
// 副作用：清除 exhausted 标记（因为我们刚刚拿到了上游确认的余额）。
func refreshAccountCreditsFromLoadCodeAssist(account *Account, resp *antigravity.LoadCodeAssistResponse) {
	if account == nil || resp == nil {
		return
	}
	if account.Extra == nil {
		account.Extra = make(map[string]any)
	}

	balance := pickGoogleOneAIBalance(resp.GetAvailableCredits())
	account.Extra[extraKeyCreditsBalance] = balance
	account.Extra[extraKeyCreditsCheckedAt] = time.Now().UTC().Format(time.RFC3339)
	// 余额已重新查询；如果有余额或耗尽标记早于本次查询，应清除。
	delete(account.Extra, extraKeyCreditsExhaustedAt)
}

// pickGoogleOneAIBalance 从 availableCredits 列表中提取 GOOGLE_ONE_AI 余额。
// creditType 为空（旧响应格式）时按整体余额累加。
func pickGoogleOneAIBalance(credits []antigravity.AvailableCredit) float64 {
	var total float64
	for _, c := range credits {
		if c.CreditType == "" || c.CreditType == antigravity.CreditTypeGoogleOneAI {
			total += c.GetAmount()
		}
	}
	return total
}

// markAccountCreditsExhausted 把账号 credits 标记为余额不足（INSUFFICIENT_G1_CREDITS_BALANCE）。
// 余额改写为 0，写入耗尽时间戳，并把更新同步到 Redis 调度快照。
func (s *AntigravityGatewayService) markAccountCreditsExhausted(ctx context.Context, prefix string, account *Account) {
	if account == nil {
		return
	}
	if account.Extra == nil {
		account.Extra = make(map[string]any)
	}
	account.Extra[extraKeyCreditsBalance] = 0.0
	account.Extra[extraKeyCreditsExhaustedAt] = time.Now().UTC().Format(time.RFC3339)
	account.Extra[extraKeyCreditsCheckedAt] = time.Now().UTC().Format(time.RFC3339)

	if s.schedulerSnapshot != nil {
		if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil {
			log.Printf("%s credits_exhausted_cache_update_failed account=%d err=%v", prefix, account.ID, err)
		}
	}
}

// readFloat 从 Account.Extra 的 any 类型里宽松读取浮点数（兼容 JSON 反序列化的 float64/string）。
func readFloat(v any) float64 {
	switch x := v.(type) {
	case float64:
		return x
	case float32:
		return float64(x)
	case int:
		return float64(x)
	case int64:
		return float64(x)
	case string:
		if x == "" {
			return 0
		}
		f, err := strconv.ParseFloat(x, 64)
		if err != nil {
			return 0
		}
		return f
	}
	return 0
}