Epicure/apps/web/lib/tiers.ts

import { db } from "@epicure/db";
import { tierDefinitions, userUsage } from "@epicure/db";
import { eq, sql } from "@epicure/db";

function currentMonth() {
  const now = new Date();
  return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`;
}

export type LimitKey = "recipe" | "aiCall" | "storage";

/** Sentinel stored in tier_definitions numeric columns to mean "no cap". */
export const UNLIMITED = -1;

export class TierLimitError extends Error {
  constructor(public readonly limit: LimitKey, public readonly tier: string) {
    super(`Tier limit reached: ${limit} (tier: ${tier})`);
    this.name = "TierLimitError";
  }
}

/**
 * Atomically checks the tier limit and increments the usage counter in a
 * single SQL statement, eliminating the TOCTOU race that existed when
 * checkTierLimit and incrementUsage were called separately.
 *
 * Throws TierLimitError if the limit has already been reached.
 * Use this instead of the separate checkTierLimit + incrementUsage pair
 * for "recipe" and "aiCall" keys.
 */
export async function checkAndIncrementTierLimit(
  userId: string,
  userTier: "free" | "pro",
  key: "recipe" | "aiCall" | "storage",
  amount = 1
): Promise<void> {
  const [tierDef] = await db
    .select()
    .from(tierDefinitions)
    .where(eq(tierDefinitions.tier, userTier));

  if (!tierDef) throw new TierLimitError(key, userTier);

  const month = currentMonth();
  const id = `${userId}-${month}`;

  if (key === "aiCall") {
    const limit = tierDef.aiCallsPerMonth;
    const cap = limit === UNLIMITED ? sql`true` : sql`user_usage.ai_calls_used < ${limit}`;
    const result = await db.execute(sql`
      INSERT INTO user_usage (id, user_id, month, ai_calls_used, recipe_count, storage_used_mb)
      VALUES (${id}, ${userId}, ${month}, 1, 0, 0)
      ON CONFLICT (user_id, month) DO UPDATE
        SET ai_calls_used = user_usage.ai_calls_used + 1
        WHERE ${cap}
      RETURNING ai_calls_used
    `);
    if (result.length === 0) {
      throw new TierLimitError("aiCall", userTier);
    }
  } else if (key === "recipe") {
    const limit = tierDef.maxRecipes;
    const cap = limit === UNLIMITED ? sql`true` : sql`user_usage.recipe_count < ${limit}`;
    const result = await db.execute(sql`
      INSERT INTO user_usage (id, user_id, month, ai_calls_used, recipe_count, storage_used_mb)
      VALUES (${id}, ${userId}, ${month}, 0, 1, 0)
      ON CONFLICT (user_id, month) DO UPDATE
        SET recipe_count = user_usage.recipe_count + 1
        WHERE ${cap}
      RETURNING recipe_count
    `);
    if (result.length === 0) {
      throw new TierLimitError("recipe", userTier);
    }
  } else {
    const limit = tierDef.storageMb;
    const cap = limit === UNLIMITED ? sql`true` : sql`user_usage.storage_used_mb + ${amount} <= ${limit}`;
    const result = await db.execute(sql`
      INSERT INTO user_usage (id, user_id, month, ai_calls_used, recipe_count, storage_used_mb)
      VALUES (${id}, ${userId}, ${month}, 0, 0, ${amount})
      ON CONFLICT (user_id, month) DO UPDATE
        SET storage_used_mb = user_usage.storage_used_mb + ${amount}
        WHERE ${cap}
      RETURNING storage_used_mb
    `);
    if (result.length === 0) {
      throw new TierLimitError("storage", userTier);
    }
  }
}

/**
 * Refunds one aiCall credit for the current month. Call this when an AI
 * request failed after the quota was already charged (e.g. provider error)
 * so users aren't billed against their limit for a call that never succeeded.
 */
export async function refundAiCall(userId: string): Promise<void> {
  const month = currentMonth();
  await db.execute(sql`
    UPDATE user_usage
    SET ai_calls_used = GREATEST(ai_calls_used - 1, 0)
    WHERE user_id = ${userId} AND month = ${month}
  `);
}

export async function incrementUsage(
  userId: string,
  key: LimitKey,
  amount = 1
): Promise<void> {
  const month = currentMonth();
  const id = `${userId}-${month}`;

  const initialValues = {
    id,
    userId,
    month,
    aiCallsUsed: key === "aiCall" ? amount : 0,
    recipeCount: key === "recipe" ? amount : 0,
    storageUsedMb: key === "storage" ? amount : 0,
  };

  const incrementSet =
    key === "aiCall"
      ? { aiCallsUsed: sql`${userUsage.aiCallsUsed} + ${amount}` }
      : key === "recipe"
        ? { recipeCount: sql`${userUsage.recipeCount} + ${amount}` }
        : { storageUsedMb: sql`${userUsage.storageUsedMb} + ${amount}` };

  await db
    .insert(userUsage)
    .values(initialValues)
    .onConflictDoUpdate({
      target: [userUsage.userId, userUsage.month],
      set: incrementSet,
    });
}