AI 功能商业化（用量/限流/计费）

// lib/ai/usage.ts
import { AI_CONFIG } from './config'
 
interface UsageData {
  promptTokens: number
  completionTokens: number
}
 
export async function recordUsage(
  tenantId: string,
  userId: string,
  model: string,
  usage: UsageData
) {
  const costs = AI_CONFIG.costs[model]
  const cost = costs
    ? (usage.promptTokens / 1000) * costs.input + (usage.completionTokens / 1000) * costs.output
    : 0
 
  await db.insert(aiUsage).values({
    tenantId,
    userId,
    model,
    inputTokens: usage.promptTokens,
    outputTokens: usage.completionTokens,
    cost: cost.toFixed(6),
    operationType: 'chat',
    createdAt: new Date(),
  })
}

// lib/ai/usage-stats.ts
 
interface UsageSummary {
  totalTokens: number
  totalCost: number
  byModel: Record<string, { tokens: number; cost: number }>
  dailyUsage: { date: string; tokens: number }[]
}
 
export async function getUsageSummary(
  tenantId: string,
  period: { start: Date; end: Date }
): Promise<UsageSummary> {
  // 总量统计
  const [totals] = await db.select({
    totalInputTokens: sql<number>`sum(input_tokens)`,
    totalOutputTokens: sql<number>`sum(output_tokens)`,
    totalCost: sql<number>`sum(cost::numeric)`,
  })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      gte(aiUsage.createdAt, period.start),
      lte(aiUsage.createdAt, period.end),
    ))
 
  // 按模型分组
  const byModel = await db.select({
    model: aiUsage.model,
    inputTokens: sql<number>`sum(input_tokens)`,
    outputTokens: sql<number>`sum(output_tokens)`,
    cost: sql<number>`sum(cost::numeric)`,
  })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      gte(aiUsage.createdAt, period.start),
      lte(aiUsage.createdAt, period.end),
    ))
    .groupBy(aiUsage.model)
 
  // 按天统计（用于图表）
  const daily = await db.select({
    date: sql<string>`date(created_at)`,
    tokens: sql<number>`sum(input_tokens + output_tokens)`,
  })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      gte(aiUsage.createdAt, period.start),
      lte(aiUsage.createdAt, period.end),
    ))
    .groupBy(sql`date(created_at)`)
    .orderBy(sql`date(created_at)`)
 
  const totalTokens = (totals.totalInputTokens || 0) + (totals.totalOutputTokens || 0)
 
  return {
    totalTokens,
    totalCost: totals.totalCost || 0,
    byModel: Object.fromEntries(
      byModel.map(m => [m.model, {
        tokens: (m.inputTokens || 0) + (m.outputTokens || 0),
        cost: m.cost || 0,
      }])
    ),
    dailyUsage: daily.map(d => ({ date: d.date, tokens: d.tokens || 0 })),
  }
}

// app/(dashboard)/settings/ai-usage/page.tsx
export default async function AIUsagePage() {
  const { membership } = await requirePermission('ai:usage:view')
 
  const now = new Date()
  const monthStart = new Date(now.getFullYear(), now.getMonth(), 1)
  const summary = await getUsageSummary(membership.tenantId, {
    start: monthStart,
    end: now,
  })
 
  const quota = getPlanQuota(membership.plan)
  const usagePercent = (summary.totalTokens / quota.monthlyTokens) * 100
 
  return (
    <div className="space-y-6">
      <h1 className="text-2xl font-bold">AI Usage</h1>
 
      {/* 配额进度 */}
      <div className="rounded-lg border p-6">
        <div className="flex items-center justify-between">
          <div>
            <p className="text-sm text-muted-foreground">Monthly Token Usage</p>
            <p className="text-2xl font-bold">
              {formatNumber(summary.totalTokens)} / {formatNumber(quota.monthlyTokens)}
            </p>
          </div>
          <div className="text-right">
            <p className="text-sm text-muted-foreground">Estimated Cost</p>
            <p className="text-2xl font-bold">${summary.totalCost.toFixed(2)}</p>
          </div>
        </div>
        <div className="mt-4 h-3 rounded-full bg-muted overflow-hidden">
          <div
            className={`h-full rounded-full transition-all ${
              usagePercent > 90 ? 'bg-red-500' : usagePercent > 70 ? 'bg-yellow-500' : 'bg-green-500'
            }`}
            style={{ width: `${Math.min(usagePercent, 100)}%` }}
          />
        </div>
        <p className="mt-2 text-sm text-muted-foreground">
          {usagePercent.toFixed(1)}% used · Resets on {formatDate(getNextResetDate())}
        </p>
      </div>
 
      {/* 按模型细分 */}
      <div className="rounded-lg border p-6">
        <h3 className="font-semibold mb-4">Usage by Model</h3>
        <table className="w-full text-sm">
          <thead>
            <tr className="border-b">
              <th className="text-left py-2">Model</th>
              <th className="text-right py-2">Tokens</th>
              <th className="text-right py-2">Cost</th>
            </tr>
          </thead>
          <tbody>
            {Object.entries(summary.byModel).map(([model, data]) => (
              <tr key={model} className="border-b">
                <td className="py-2 font-mono text-xs">{model}</td>
                <td className="py-2 text-right">{formatNumber(data.tokens)}</td>
                <td className="py-2 text-right">${data.cost.toFixed(4)}</td>
              </tr>
            ))}
          </tbody>
        </table>
      </div>
    </div>
  )
}

// lib/ai/quota.ts
 
interface PlanQuota {
  monthlyTokens: number
  maxTokensPerRequest: number
  allowedModels: string[]
  maxKnowledgeBases: number
  maxDocumentsPerKB: number
  maxAgentSteps: number
}
 
const PLAN_QUOTAS: Record<string, PlanQuota> = {
  free: {
    monthlyTokens: 50_000,
    maxTokensPerRequest: 2000,
    allowedModels: ['gpt-4o-mini'],
    maxKnowledgeBases: 1,
    maxDocumentsPerKB: 10,
    maxAgentSteps: 2,
  },
  pro: {
    monthlyTokens: 1_000_000,
    maxTokensPerRequest: 8000,
    allowedModels: ['gpt-4o-mini', 'gpt-4o'],
    maxKnowledgeBases: 5,
    maxDocumentsPerKB: 100,
    maxAgentSteps: 5,
  },
  business: {
    monthlyTokens: 10_000_000,
    maxTokensPerRequest: 32000,
    allowedModels: ['gpt-4o-mini', 'gpt-4o', 'claude-3-5-sonnet'],
    maxKnowledgeBases: 20,
    maxDocumentsPerKB: 1000,
    maxAgentSteps: 10,
  },
}
 
export function getPlanQuota(plan: string): PlanQuota {
  return PLAN_QUOTAS[plan] || PLAN_QUOTAS.free
}

export async function checkAIQuota(tenantId: string, userId: string) {
  // 获取当前套餐
  const [tenant] = await db.select().from(tenants).where(eq(tenants.id, tenantId))
  const quota = getPlanQuota(tenant.plan)
 
  // 查询当月已用量
  const monthStart = new Date()
  monthStart.setDate(1)
  monthStart.setHours(0, 0, 0, 0)
 
  const [usage] = await db.select({
    totalTokens: sql<number>`sum(input_tokens + output_tokens)`,
  })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      gte(aiUsage.createdAt, monthStart),
    ))
 
  const usedTokens = usage.totalTokens || 0
 
  if (usedTokens >= quota.monthlyTokens) {
    throw new AIQuotaExceededError(
      `Monthly AI quota exceeded. Used ${formatNumber(usedTokens)} of ${formatNumber(quota.monthlyTokens)} tokens.`,
      { usedTokens, quota: quota.monthlyTokens, plan: tenant.plan }
    )
  }
 
  // 返回剩余额度，用于 maxTokens 限制
  const remainingTokens = quota.monthlyTokens - usedTokens
  return {
    remainingTokens,
    maxTokensPerRequest: Math.min(quota.maxTokensPerRequest, remainingTokens),
    allowedModels: quota.allowedModels,
  }
}
 
class AIQuotaExceededError extends Error {
  constructor(message: string, public details: any) {
    super(message)
    this.name = 'AIQuotaExceededError'
  }
}

// lib/ai/quota-cache.ts
import { Redis } from '@upstash/redis'
 
const redis = new Redis({
  url: process.env.UPSTASH_REDIS_URL!,
  token: process.env.UPSTASH_REDIS_TOKEN!,
})
 
function getQuotaKey(tenantId: string): string {
  const month = new Date().toISOString().slice(0, 7) // '2024-01'
  return `ai:quota:${tenantId}:${month}`
}
 
export async function getCachedUsage(tenantId: string): Promise<number> {
  const key = getQuotaKey(tenantId)
  const cached = await redis.get<number>(key)
 
  if (cached !== null) return cached
 
  // 缓存未命中，从数据库加载
  const usage = await getMonthlyUsageFromDB(tenantId)
  await redis.set(key, usage, { ex: 3600 }) // 1 小时过期
  return usage
}
 
export async function incrementCachedUsage(tenantId: string, tokens: number) {
  const key = getQuotaKey(tenantId)
  await redis.incrby(key, tokens)
}

export function getModelByPlan(
  requestedModel: string,
  plan: string
): ReturnType<typeof openai> {
  const quota = getPlanQuota(plan)
 
  if (!quota.allowedModels.includes(requestedModel)) {
    // 降级到套餐允许的最佳模型
    const fallback = quota.allowedModels[quota.allowedModels.length - 1]
    return openai(fallback)
  }
 
  return openai(requestedModel)
}

// lib/ai/rate-limit.ts
import { Ratelimit } from '@upstash/ratelimit'
import { Redis } from '@upstash/redis'
 
const redis = new Redis({
  url: process.env.UPSTASH_REDIS_URL!,
  token: process.env.UPSTASH_REDIS_TOKEN!,
})
 
// 每用户限流
const userLimiter = new Ratelimit({
  redis,
  limiter: Ratelimit.slidingWindow(20, '1 m'),
  prefix: 'ai:rl:user',
})
 
// 每租户限流
const tenantLimiter = new Ratelimit({
  redis,
  limiter: Ratelimit.slidingWindow(100, '1 m'),
  prefix: 'ai:rl:tenant',
})
 
export async function checkRateLimit(tenantId: string, userId: string) {
  const [userResult, tenantResult] = await Promise.all([
    userLimiter.limit(userId),
    tenantLimiter.limit(tenantId),
  ])
 
  if (!userResult.success) {
    throw new Error(`Rate limit exceeded. Try again in ${Math.ceil(userResult.reset - Date.now()) / 1000}s`)
  }
 
  if (!tenantResult.success) {
    throw new Error('Organization rate limit exceeded. Please try again shortly.')
  }
}

// lib/ai/concurrency.ts
 
export async function acquireLock(userId: string, ttlMs: number = 60000): Promise<boolean> {
  const key = `ai:lock:${userId}`
  const acquired = await redis.set(key, '1', { nx: true, px: ttlMs })
  return acquired === 'OK'
}
 
export async function releaseLock(userId: string) {
  await redis.del(`ai:lock:${userId}`)
}
 
// 在 API Route 中使用
export async function POST(req: Request) {
  const session = await getSession()
  
  if (!await acquireLock(session!.user.id)) {
    return new Response('Another AI request is still processing', { status: 429 })
  }
 
  try {
    // ... 处理 AI 请求
  } finally {
    await releaseLock(session!.user.id)
  }
}

export async function getEffectiveModel(
  requestedModel: string,
  tenantId: string,
  plan: string
): Promise<{ model: string; isDowngraded: boolean }> {
  const remaining = await getCachedUsage(tenantId)
  const quota = getPlanQuota(plan)
 
  if (remaining < quota.monthlyTokens) {
    // 额度充足，使用请求的模型
    return { model: requestedModel, isDowngraded: false }
  }
 
  // 检查是否有额外购买的额度
  const extraQuota = await getExtraQuota(tenantId)
  if (extraQuota > 0) {
    return { model: requestedModel, isDowngraded: false }
  }
 
  // 降级到最便宜的模型
  return { model: 'gpt-4o-mini', isDowngraded: true }
}

// lib/actions/ai-credits.ts
'use server'
 
export async function purchaseAICredits(tokenAmount: number) {
  const { session, membership } = await requirePermission('ai:usage:view')
 
  // 价格：$10 / 1M tokens
  const PRICE_PER_MILLION = 10
  const amount = Math.ceil((tokenAmount / 1_000_000) * PRICE_PER_MILLION * 100) // cents
 
  // 创建 Stripe Checkout Session
  const checkoutSession = await stripe.checkout.sessions.create({
    mode: 'payment',
    customer: membership.stripeCustomerId,
    line_items: [{
      price_data: {
        currency: 'usd',
        product_data: {
          name: `${formatNumber(tokenAmount)} AI Tokens`,
          description: 'Additional AI tokens for your organization',
        },
        unit_amount: amount,
      },
      quantity: 1,
    }],
    metadata: {
      tenantId: membership.tenantId,
      tokenAmount: String(tokenAmount),
      type: 'ai_credits',
    },
    success_url: `${process.env.NEXT_PUBLIC_APP_URL}/settings/ai-usage?purchased=true`,
    cancel_url: `${process.env.NEXT_PUBLIC_APP_URL}/settings/ai-usage`,
  })
 
  return { url: checkoutSession.url }
}

// 在 Stripe Webhook 中处理 ai_credits 类型
if (session.metadata?.type === 'ai_credits') {
  const tenantId = session.metadata.tenantId
  const tokenAmount = parseInt(session.metadata.tokenAmount)
 
  await db.insert(aiCredits).values({
    tenantId,
    tokens: tokenAmount,
    expiresAt: new Date(Date.now() + 90 * 24 * 60 * 60 * 1000), // 90 天有效
  })
 
  // 清除配额缓存
  await redis.del(getQuotaKey(tenantId))
}

// lib/ai/cost-alert.ts
 
export async function checkCostAnomaly(tenantId: string) {
  const today = new Date()
  const todayStart = new Date(today.setHours(0, 0, 0, 0))
 
  // 今日成本
  const [todayCost] = await db.select({
    cost: sql<number>`sum(cost::numeric)`,
  })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      gte(aiUsage.createdAt, todayStart),
    ))
 
  // 过去 7 天日均成本
  const weekAgo = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000)
  const [avgCost] = await db.select({
    avgCost: sql<number>`sum(cost::numeric) / 7`,
  })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      gte(aiUsage.createdAt, weekAgo),
      lte(aiUsage.createdAt, todayStart),
    ))
 
  const dailyAvg = avgCost.avgCost || 0
  const todayTotal = todayCost.cost || 0
 
  // 如果今日成本超过日均的 3 倍，触发告警
  if (dailyAvg > 0 && todayTotal > dailyAvg * 3) {
    await sendCostAlert(tenantId, {
      todayCost: todayTotal,
      dailyAverage: dailyAvg,
      ratio: todayTotal / dailyAvg,
    })
  }
}

// lib/actions/admin-ai-stats.ts
 
export async function getAICostOverview() {
  const monthStart = new Date()
  monthStart.setDate(1)
  monthStart.setHours(0, 0, 0, 0)
 
  // 全平台本月成本
  const [overview] = await db.select({
    totalCost: sql<number>`sum(cost::numeric)`,
    totalTokens: sql<number>`sum(input_tokens + output_tokens)`,
    uniqueUsers: sql<number>`count(distinct user_id)`,
    totalRequests: sql<number>`count(*)`,
  })
    .from(aiUsage)
    .where(gte(aiUsage.createdAt, monthStart))
 
  // 按租户的 Top 10 消耗
  const topTenants = await db.select({
    tenantId: aiUsage.tenantId,
    cost: sql<number>`sum(cost::numeric)`,
    tokens: sql<number>`sum(input_tokens + output_tokens)`,
  })
    .from(aiUsage)
    .where(gte(aiUsage.createdAt, monthStart))
    .groupBy(aiUsage.tenantId)
    .orderBy(desc(sql`sum(cost::numeric)`))
    .limit(10)
 
  return { overview, topTenants }
}

// lib/ai/circuit-breaker.ts
 
const DAILY_BUDGET = parseFloat(process.env.AI_DAILY_BUDGET || '100')  // $100/天
 
export async function checkGlobalBudget(): Promise<boolean> {
  const key = `ai:daily-cost:${new Date().toISOString().slice(0, 10)}`
  const currentCost = await redis.get<number>(key) || 0
 
  if (currentCost >= DAILY_BUDGET) {
    // 触发熔断：只允许 mini 模型
    return false
  }
 
  return true
}
 
export async function incrementGlobalCost(cost: number) {
  const key = `ai:daily-cost:${new Date().toISOString().slice(0, 10)}`
  await redis.incrbyfloat(key, cost)
  await redis.expire(key, 86400 * 2)  // 2 天过期
}