Epicure/apps/web/lib/tiers.ts

import { db } from "@epicure/db";
import { tierDefinitions, users, recipes, recipePhotos, ratings } 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";
  }
}

/** Either the top-level db client or a transaction handle — both expose the
 * same query-builder surface, so live-derivation helpers and the limit check
 * can run against whichever one the caller is holding a lock in. */
type DbOrTx = typeof db | Parameters<Parameters<typeof db.transaction>[0]>[0];

/** Live count of a user's recipes — lifetime total, not a monthly counter. */
export async function getRecipeCount(userId: string, dbOrTx: DbOrTx = db): Promise<number> {
  const [row] = await dbOrTx
    .select({ n: sql<number>`count(*)::int` })
    .from(recipes)
    .where(eq(recipes.authorId, userId));
  return row?.n ?? 0;
}

/** Live sum of a user's storage usage (recipe photos + review photos + avatar) — lifetime total, not a monthly counter. */
export async function getStorageUsedMb(userId: string, dbOrTx: DbOrTx = db): Promise<number> {
  const [[photoRow], [reviewRow], [userRow]] = await Promise.all([
    dbOrTx
      .select({ total: sql<number>`coalesce(sum(${recipePhotos.sizeMb}), 0)::int` })
      .from(recipePhotos)
      .innerJoin(recipes, eq(recipePhotos.recipeId, recipes.id))
      .where(eq(recipes.authorId, userId)),
    dbOrTx
      .select({ total: sql<number>`coalesce(sum(${ratings.photoSizeMb}), 0)::int` })
      .from(ratings)
      .where(eq(ratings.userId, userId)),
    dbOrTx.select({ avatarSizeMb: users.avatarSizeMb }).from(users).where(eq(users.id, userId)),
  ]);
  return (photoRow?.total ?? 0) + (reviewRow?.total ?? 0) + (userRow?.avatarSizeMb ?? 0);
}

async function checkLiveLimit(
  dbOrTx: DbOrTx,
  userId: string,
  fallbackTier: "free" | "pro" | "family",
  key: "recipe" | "storage",
  amount: number
): Promise<void> {
  const [dbUser] = await dbOrTx.select({ tier: users.tier }).from(users).where(eq(users.id, userId));
  const userTier = (dbUser?.tier as "free" | "pro" | "family" | undefined) ?? fallbackTier;

  const [tierDef] = await dbOrTx.select().from(tierDefinitions).where(eq(tierDefinitions.tier, userTier));
  if (!tierDef) throw new TierLimitError(key, userTier);

  const limit = key === "recipe" ? tierDef.maxRecipes : tierDef.storageMb;
  if (limit === UNLIMITED) return;

  const current = key === "recipe" ? await getRecipeCount(userId, dbOrTx) : await getStorageUsedMb(userId, dbOrTx);
  if (current + amount > limit) throw new TierLimitError(key, userTier);
}

/**
 * Checks the tier limit for the given key, throwing TierLimitError if it's
 * already been reached (or would be exceeded by `amount`).
 *
 * The caller's `userTier` is never trusted directly — it comes from the
 * session's 5-minute cookieCache (see lib/auth/server.ts), so a just-downgraded
 * user would otherwise keep the old tier's caps for up to 5 minutes. The
 * current tier is always re-read from the DB here.
 *
 * "aiCall" is the only key that's monthly — it atomically checks-and-increments
 * a per-month counter in a single SQL statement to avoid a TOCTOU race.
 * "recipe" and "storage" are lifetime totals derived live from real data
 * (recipes/photos/avatar rows), so they're pure checks with no counter to
 * increment — deleting a recipe/photo/avatar is itself the "decrement".
 *
 * This plain (unlocked) check is a fast pre-flight only — two concurrent
 * calls can both read the pre-write count and both pass. Anywhere the actual
 * write happens in the same request, use checkTierLimitInTransaction instead
 * so the check and the write are atomic together.
 */
export async function checkAndIncrementTierLimit(
  userId: string,
  fallbackTier: "free" | "pro" | "family",
  key: "recipe" | "aiCall" | "storage",
  amount = 1
): Promise<void> {
  if (key === "aiCall") {
    const [dbUser] = await db.select({ tier: users.tier }).from(users).where(eq(users.id, userId));
    const userTier = (dbUser?.tier as "free" | "pro" | "family" | undefined) ?? fallbackTier;

    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}`;
    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)
      VALUES (${id}, ${userId}, ${month}, 1)
      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);
    }
    return;
  }

  await checkLiveLimit(db, userId, fallbackTier, key, amount);
}

/**
 * Same check as checkAndIncrementTierLimit's "recipe"/"storage" branches, but
 * takes a transaction and holds a per-user Postgres advisory lock for its
 * duration — so a concurrent call for the same user blocks until this one
 * commits (or rolls back), instead of racing to read the same pre-write
 * count. The lock auto-releases at transaction end (pg_advisory_xact_lock),
 * so there's no separate unlock step and no risk of a leaked lock.
 *
 * Call this as the first statement inside the same db.transaction that
 * performs the write the check is gating — checking and inserting in
 * different transactions (or different requests, like a presign-time check
 * for a row written later) leaves the same TOCTOU gap this closes.
 */
export async function checkTierLimitInTransaction(
  tx: DbOrTx,
  userId: string,
  fallbackTier: "free" | "pro" | "family",
  key: "recipe" | "storage",
  amount = 1
): Promise<void> {
  await tx.execute(sql`SELECT pg_advisory_xact_lock(hashtext(${userId}))`);
  await checkLiveLimit(tx, userId, fallbackTier, key, amount);
}

/**
 * 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}
  `);
}