Momento/memento-note/lib/ai-quota.ts

import { willUseByokForLane } from '@/lib/ai/provider-for-user'
import type { AiFeatureLane } from '@/lib/ai/router'
import { getSystemConfig } from '@/lib/config'
import {
  checkSessionEntitlementOrThrow,
  QuotaExceededError,
  QuotaServiceUnavailableError,
} from '@/lib/entitlements'
import {
  releaseCredits,
  reserveCreditsOrThrow,
  resolveCreditCost,
} from '@/lib/credits'
import type { FeatureName } from '@/lib/quota-utils'
import { emitAiUsageChanged } from '@/lib/ai-usage-sync'

export type ReserveAiUsageOptions = {
  lane?: AiFeatureLane
  /** Defaults to userId — host-pays brainstorm bills this account. */
  billingUserId?: string
  /** Units / crédits (default 1). Multi-step features (slides) pass amount > 1. */
  amount?: number
  /** Pour coût slides précis (1+N). */
  slideCount?: number | null
  metadata?: Record<string, unknown>
}

/** Reserve crédits IA ; saute si BYOK pour le lane. */
export async function reserveAiUsageOrThrow(
  userId: string,
  feature: FeatureName,
  options?: ReserveAiUsageOptions,
): Promise<{ usedByok: boolean; billedUserId: string; amount: number }> {
  const billedUserId = options?.billingUserId ?? userId
  const lane = options?.lane ?? 'chat'
  const cost = resolveCreditCost(feature, {
    amount: options?.amount,
    slideCount: options?.slideCount,
  })
  const config = await getSystemConfig()
  const { usedByok } = await willUseByokForLane(lane, config, billedUserId)
  if (!usedByok) {
    await reserveCreditsOrThrow(billedUserId, cost, {
      feature,
      metadata: options?.metadata,
    })
  }
  return { usedByok, billedUserId, amount: cost }
}

/** Host-pays: reserve on session owner with guest metadata on 402. */
export async function reserveSessionAiUsageOrThrow(
  billingOwnerId: string,
  triggeredByUserId: string,
  isGuestActor: boolean,
  feature: FeatureName,
  options?: { lane?: AiFeatureLane; amount?: number; slideCount?: number | null },
): Promise<{ usedByok: boolean; amount: number }> {
  const lane = options?.lane ?? 'chat'
  const cost = resolveCreditCost(feature, {
    amount: options?.amount,
    slideCount: options?.slideCount,
  })
  const config = await getSystemConfig()
  const { usedByok } = await willUseByokForLane(lane, config, billingOwnerId)
  if (!usedByok) {
    try {
      await reserveCreditsOrThrow(billingOwnerId, cost, { feature })
    } catch (err) {
      if (err instanceof QuotaExceededError) {
        err.billingOwnerId = billingOwnerId
        err.triggeredByUserId = triggeredByUserId
        err.isGuestActor = isGuestActor
        // Also keep session path for legacy callers
        try {
          await checkSessionEntitlementOrThrow(
            billingOwnerId,
            triggeredByUserId,
            isGuestActor,
            feature,
          )
        } catch {
          /* already QuotaExceeded from credits */
        }
      }
      throw err
    }
  }
  return { usedByok, amount: cost }
}

/** Run fn under crédits ; libère si fn throw (sauf BYOK). */
export async function withAiQuota<T>(
  userId: string,
  feature: FeatureName,
  fn: () => Promise<T>,
  options?: ReserveAiUsageOptions,
): Promise<T> {
  const { usedByok, billedUserId, amount } = await reserveAiUsageOrThrow(userId, feature, options)
  try {
    const result = await fn()
    if (!usedByok) {
      emitAiUsageChanged()
    }
    return result
  } catch (err) {
    if (!usedByok && amount > 0) {
      await releaseCredits(billedUserId, amount, { feature })
    }
    throw err
  }
}

/** Host-pays collaborative flows (brainstorm) with rollback on failure. */
export async function withSessionAiQuota<T>(
  billingOwnerId: string,
  triggeredByUserId: string,
  isGuestActor: boolean,
  feature: FeatureName,
  fn: () => Promise<T>,
  options?: { lane?: AiFeatureLane; amount?: number; slideCount?: number | null },
): Promise<T> {
  const { usedByok, amount } = await reserveSessionAiUsageOrThrow(
    billingOwnerId,
    triggeredByUserId,
    isGuestActor,
    feature,
    options,
  )
  try {
    const result = await fn()
    if (!usedByok) {
      emitAiUsageChanged()
    }
    return result
  } catch (err) {
    if (!usedByok && amount > 0) {
      await releaseCredits(billingOwnerId, amount, { feature })
    }
    throw err
  }
}

export function isQuotaError(err: unknown): err is QuotaExceededError {
  return err instanceof QuotaExceededError
}

export function quotaExceededJson(err: QuotaExceededError) {
  return err.toJSON()
}

export function quotaErrorStatus() {
  return 402
}

/** Standard Next.js / Response handlers for quota errors. */
export function handleQuotaHttpError(err: unknown): Response | null {
  if (err instanceof QuotaExceededError) {
    return Response.json(err.toJSON(), { status: 402 })
  }
  if (err instanceof QuotaServiceUnavailableError) {
    return Response.json({ error: err.code }, { status: 503 })
  }
  return null
}

export function notifyAiUsageChanged(): void {
  emitAiUsageChanged()
}