claude-relay-service/src/services/balanceProviders/geminiBalanceProvider.js

const BaseBalanceProvider = require('./baseBalanceProvider')
const antigravityClient = require('../antigravityClient')
const geminiAccountService = require('../geminiAccountService')

const OAUTH_PROVIDER_ANTIGRAVITY = 'antigravity'

function clamp01(value) {
  if (typeof value !== 'number' || !Number.isFinite(value)) {
    return null
  }
  if (value < 0) {
    return 0
  }
  if (value > 1) {
    return 1
  }
  return value
}

function round2(value) {
  if (typeof value !== 'number' || !Number.isFinite(value)) {
    return null
  }
  return Math.round(value * 100) / 100
}

function normalizeQuotaCategory(displayName, modelId) {
  const name = String(displayName || '')
  const id = String(modelId || '')

  if (name.includes('Gemini') && name.includes('Pro')) {
    return 'Gemini Pro'
  }
  if (name.includes('Gemini') && name.includes('Flash')) {
    return 'Gemini Flash'
  }
  if (name.includes('Gemini') && name.toLowerCase().includes('image')) {
    return 'Gemini Image'
  }

  if (name.includes('Claude') || name.includes('GPT-OSS')) {
    return 'Claude'
  }

  if (id.startsWith('gemini-3-pro-') || id.startsWith('gemini-2.5-pro')) {
    return 'Gemini Pro'
  }
  if (id.startsWith('gemini-3-flash') || id.startsWith('gemini-2.5-flash')) {
    return 'Gemini Flash'
  }
  if (id.includes('image')) {
    return 'Gemini Image'
  }
  if (id.includes('claude') || id.includes('gpt-oss')) {
    return 'Claude'
  }

  return name || id || 'Unknown'
}

function buildAntigravityQuota(modelsResponse) {
  const models = modelsResponse && typeof modelsResponse === 'object' ? modelsResponse.models : null

  if (!models || typeof models !== 'object') {
    return null
  }

  const parseRemainingFraction = (quotaInfo) => {
    if (!quotaInfo || typeof quotaInfo !== 'object') {
      return null
    }

    const raw =
      quotaInfo.remainingFraction ??
      quotaInfo.remaining_fraction ??
      quotaInfo.remaining ??
      undefined

    const num = typeof raw === 'number' ? raw : typeof raw === 'string' ? Number(raw) : NaN
    if (!Number.isFinite(num)) {
      return null
    }

    return clamp01(num)
  }

  const allowedCategories = new Set(['Gemini Pro', 'Claude', 'Gemini Flash', 'Gemini Image'])
  const fixedOrder = ['Gemini Pro', 'Claude', 'Gemini Flash', 'Gemini Image']

  const categoryMap = new Map()

  for (const [modelId, modelDataRaw] of Object.entries(models)) {
    if (!modelDataRaw || typeof modelDataRaw !== 'object') {
      continue
    }

    const displayName = modelDataRaw.displayName || modelDataRaw.display_name || modelId
    const quotaInfo = modelDataRaw.quotaInfo || modelDataRaw.quota_info || null

    const remainingFraction = parseRemainingFraction(quotaInfo)
    if (remainingFraction === null) {
      continue
    }

    const remainingPercent = round2(remainingFraction * 100)
    const usedPercent = round2(100 - remainingPercent)
    const resetAt = quotaInfo?.resetTime || quotaInfo?.reset_time || null

    const category = normalizeQuotaCategory(displayName, modelId)
    if (!allowedCategories.has(category)) {
      continue
    }
    const entry = {
      category,
      modelId,
      displayName: String(displayName || modelId || category),
      remainingPercent,
      usedPercent,
      resetAt: typeof resetAt === 'string' && resetAt.trim() ? resetAt : null
    }

    const existing = categoryMap.get(category)
    if (!existing || entry.remainingPercent < existing.remainingPercent) {
      categoryMap.set(category, entry)
    }
  }

  const buckets = fixedOrder.map((category) => {
    const existing = categoryMap.get(category) || null
    if (existing) {
      return existing
    }
    return {
      category,
      modelId: '',
      displayName: category,
      remainingPercent: null,
      usedPercent: null,
      resetAt: null
    }
  })

  if (buckets.length === 0) {
    return null
  }

  const critical = buckets
    .filter((item) => item.remainingPercent !== null)
    .reduce((min, item) => {
      if (!min) {
        return item
      }
      return (item.remainingPercent ?? 0) < (min.remainingPercent ?? 0) ? item : min
    }, null)

  if (!critical) {
    return null
  }

  return {
    balance: null,
    currency: 'USD',
    quota: {
      type: 'antigravity',
      total: 100,
      used: critical.usedPercent,
      remaining: critical.remainingPercent,
      percentage: critical.usedPercent,
      resetAt: critical.resetAt,
      buckets: buckets.map((item) => ({
        category: item.category,
        remaining: item.remainingPercent,
        used: item.usedPercent,
        percentage: item.usedPercent,
        resetAt: item.resetAt
      }))
    },
    queryMethod: 'api',
    rawData: {
      modelsCount: Object.keys(models).length,
      bucketCount: buckets.length
    }
  }
}

class GeminiBalanceProvider extends BaseBalanceProvider {
  constructor() {
    super('gemini')
  }

  async queryBalance(account) {
    const oauthProvider = account?.oauthProvider
    if (oauthProvider !== OAUTH_PROVIDER_ANTIGRAVITY) {
      if (account && Object.prototype.hasOwnProperty.call(account, 'dailyQuota')) {
        return this.readQuotaFromFields(account)
      }
      return { balance: null, currency: 'USD', queryMethod: 'local' }
    }

    const accessToken = String(account?.accessToken || '').trim()
    const refreshToken = String(account?.refreshToken || '').trim()
    const proxyConfig = account?.proxyConfig || account?.proxy || null

    if (!accessToken) {
      throw new Error('Antigravity 账户缺少 accessToken')
    }

    const fetch = async (token) =>
      await antigravityClient.fetchAvailableModels({
        accessToken: token,
        proxyConfig
      })

    let data
    try {
      data = await fetch(accessToken)
    } catch (error) {
      const status = error?.response?.status
      if ((status === 401 || status === 403) && refreshToken) {
        const refreshed = await geminiAccountService.refreshAccessToken(
          refreshToken,
          proxyConfig,
          OAUTH_PROVIDER_ANTIGRAVITY
        )
        const nextToken = String(refreshed?.access_token || '').trim()
        if (!nextToken) {
          throw error
        }
        data = await fetch(nextToken)
      } else {
        throw error
      }
    }

    const mapped = buildAntigravityQuota(data)
    if (!mapped) {
      return {
        balance: null,
        currency: 'USD',
        quota: null,
        queryMethod: 'api',
        rawData: data || null
      }
    }

    return mapped
  }
}

module.exports = GeminiBalanceProvider