songs2vid/lib/quota.ts

import { Plan } from "@prisma/client";
import {
  CreditInsufficientError,
  applyMonthlyCreditRenewal,
  deductUserCredit,
  refundUserCredit,
} from "./billing";
import { FREE_EXTRA_CREDITS_MAX, MAX_CREDIT_CAP, PRO_UPGRADE_REQUIRED_SUFFIX, monthlyCreditsForPlan } from "./credits";
import { prisma } from "./db";
import { hasProFeatures, isSelfHostedEdition } from "./edition";
import { getNextMonthlyQuotaReset, getPlanLimits } from "./plans";
import { watermarkFreeRendersRemaining } from "./watermark-policy";

export function formatQuotaResetCountdown(resetsAt: Date | string): string {
  const target = typeof resetsAt === "string" ? new Date(resetsAt) : resetsAt;
  const ms = Math.max(0, target.getTime() - Date.now());
  const totalSec = Math.ceil(ms / 1000);
  const hours = Math.floor(totalSec / 3600);
  const minutes = Math.floor((totalSec % 3600) / 60);
  const seconds = totalSec % 60;
  return `${hours}h ${minutes}m ${seconds}s`;
}

export function formatQuotaResetDisplay(plan: Plan, resetsAt: Date | string): string {
  const target = typeof resetsAt === "string" ? new Date(resetsAt) : resetsAt;
  return target.toLocaleDateString(undefined, {
    month: "long",
    day: "numeric",
    year: "numeric",
  });
}

function getQuotaExceededMessage(
  plan: Plan,
  resetsAt: Date,
  extraCredits: number,
): string {
  const resetLabel = formatQuotaResetDisplay(plan, resetsAt);
  if (plan === "PREMIUM") {
    return extraCredits > 0
      ? `Monthly Pro quota exhausted. You still have ${extraCredits} extra credits.`
      : `Quota exceeded. Your Pro monthly credits reset on ${resetLabel}.`;
  }
  if (extraCredits > 0) {
    return `Monthly free quota exhausted. You still have ${extraCredits} extra credits.`;
  }
  return (
    `You've used your free monthly credits (10) and any extra top-ups (max ${FREE_EXTRA_CREDITS_MAX}). ` +
    `Upgrade to the Pro monthly plan to continue uploading.` +
    PRO_UPGRADE_REQUIRED_SUFFIX
  );
}

/** Reset cycle when due: rollover unused into extras, add new monthly, trim to MAX_CREDIT_CAP. */
export async function ensureQuotaReset(userId: string) {
  const user = await prisma.user.findUniqueOrThrow({ where: { id: userId } });
  const now = new Date();
  const expectedMonthly = monthlyCreditsForPlan(user.plan);

  if (now >= user.quotaResetAt) {
    return applyMonthlyCreditRenewal(userId, {
      plan: user.plan,
      newMonthlyCredits: expectedMonthly,
      quotaResetAt: getNextMonthlyQuotaReset(now),
    });
  }

  // Keep monthlyCredits in sync if plan was changed mid-cycle without reset
  if (user.monthlyCredits !== expectedMonthly && user.videosUsed === 0 && user.extraCredits === 0) {
    return prisma.user.update({
      where: { id: userId },
      data: { monthlyCredits: expectedMonthly },
    });
  }

  return user;
}

export async function getQuotaInfo(userId: string) {
  const user = await ensureQuotaReset(userId);
  const limits = getPlanLimits(user.plan);
  const limit = isSelfHostedEdition()
    ? limits.monthlyQuota
    : user.monthlyCredits + (user.plan === "PREMIUM" ? user.bonusQuota : 0);
  const remaining = Math.max(0, limit - user.videosUsed);
  const extraCredits = isSelfHostedEdition() ? 0 : user.extraCredits;

  return {
    used: user.videosUsed,
    creditsUsed: user.videosUsed,
    limit,
    monthlyCredits: user.monthlyCredits,
    baseLimit: limits.monthlyQuota,
    bonusQuota: isSelfHostedEdition() ? 0 : user.bonusQuota,
    remaining,
    /** @deprecated use extraCredits */
    videoCredits: extraCredits,
    extraCredits,
    creditBalanceMax: MAX_CREDIT_CAP,
    requiresProUpgrade:
      !isSelfHostedEdition() &&
      user.plan === "FREE" &&
      remaining === 0 &&
      extraCredits === 0,
    freeTopUpPurchased: user.freeTopUpPurchased,
    totalAvailable: remaining + extraCredits,
    resetsAt: user.quotaResetAt.toISOString(),
    resetsIn: formatQuotaResetDisplay(user.plan, user.quotaResetAt),
    plan: user.plan,
    subscriptionStatus: hasProFeatures(user.plan) ? ("pro" as const) : ("free" as const),
    createdVideoCount: user.createdVideoCount,
    watermarkFreeRemaining: hasProFeatures(user.plan)
      ? null
      : watermarkFreeRendersRemaining(user.createdVideoCount),
    maxBatchSize: limits.maxBatchSize,
    watermarkOptional: limits.watermarkOptional,
    customWatermark: limits.customWatermark,
    perItemImages: limits.perItemImages,
    artTrackLayouts: limits.artTrackLayouts,
    maxResolutionHeight: limits.maxResolutionHeight,
    selfHosted: isSelfHostedEdition(),
  };
}

export type ReservationSplit = {
  fromQuota: number;
  fromCredits: number;
};

/**
 * Monthly first, then never-expiring extras  for both FREE and PREMIUM.
 */
export function planReservation(
  _plan: Plan,
  remainingQuota: number,
  extraCredits: number,
  count: number,
): ReservationSplit | null {
  if (count <= 0) return { fromQuota: 0, fromCredits: 0 };
  const fromQuota = Math.min(count, Math.max(0, remainingQuota));
  const fromCredits = count - fromQuota;
  if (fromCredits > extraCredits) return null;
  return { fromQuota, fromCredits };
}

export async function checkQuota(userId: string, requestedCount: number) {
  if (isSelfHostedEdition()) {
    const info = await getQuotaInfo(userId);
    if (requestedCount > info.maxBatchSize) {
      return {
        ok: false as const,
        error: `Batch limit exceeded. Max ${info.maxBatchSize} files per batch.`,
        used: info.used,
        limit: info.limit,
        remaining: info.remaining,
        videoCredits: 0,
        extraCredits: 0,
        resetsAt: info.resetsAt,
      };
    }
    return {
      ok: true as const,
      ...info,
      reservation: { fromQuota: requestedCount, fromCredits: 0 },
    };
  }

  const user = await ensureQuotaReset(userId);
  const limits = getPlanLimits(user.plan);
  const limit = user.monthlyCredits + (user.plan === "PREMIUM" ? user.bonusQuota : 0);
  const remaining = Math.max(0, limit - user.videosUsed);

  if (requestedCount > limits.maxBatchSize) {
    return {
      ok: false as const,
      error: `Batch limit exceeded. Your plan allows up to ${limits.maxBatchSize} files per batch.`,
      used: user.videosUsed,
      limit,
      remaining,
      videoCredits: user.extraCredits,
      extraCredits: user.extraCredits,
      resetsAt: user.quotaResetAt.toISOString(),
    };
  }

  const split = planReservation(user.plan, remaining, user.extraCredits, requestedCount);
  if (!split) {
    return {
      ok: false as const,
      error: getQuotaExceededMessage(user.plan, user.quotaResetAt, user.extraCredits),
      used: user.videosUsed,
      limit,
      remaining,
      videoCredits: user.extraCredits,
      extraCredits: user.extraCredits,
      resetsAt: user.quotaResetAt.toISOString(),
    };
  }

  const info = await getQuotaInfo(userId);
  return { ok: true as const, ...info, reservation: split };
}

/**
 * Atomically reserve monthly credits first, then extras.
 */
export async function reserveQuota(userId: string, count: number): Promise<ReservationSplit> {
  if (count <= 0) return { fromQuota: 0, fromCredits: 0 };

  if (isSelfHostedEdition()) {
    const limits = getPlanLimits("FREE");
    if (count > limits.maxBatchSize) {
      throw new Error(`Batch limit exceeded. Max ${limits.maxBatchSize} files per batch.`);
    }
    await ensureQuotaReset(userId);
    await prisma.user.update({
      where: { id: userId },
      data: { videosUsed: { increment: count } },
    });
    return { fromQuota: count, fromCredits: 0 };
  }

  await ensureQuotaReset(userId);

  const limits = getPlanLimits(
    (await prisma.user.findUniqueOrThrow({ where: { id: userId } })).plan,
  );
  if (count > limits.maxBatchSize) {
    throw new Error(
      `Batch limit exceeded. Your plan allows up to ${limits.maxBatchSize} files per batch.`,
    );
  }

  try {
    const result = await deductUserCredit(userId, count);
    return { fromQuota: result.fromMonthly, fromCredits: result.fromExtra };
  } catch (err) {
    if (err instanceof CreditInsufficientError) {
      const user = await prisma.user.findUniqueOrThrow({ where: { id: userId } });
      throw new Error(
        getQuotaExceededMessage(user.plan, user.quotaResetAt, user.extraCredits),
      );
    }
    throw err;
  }
}

export async function releaseQuota(userId: string, count: number) {
  if (count <= 0) return;
  await refundUserCredit(userId, count, 0);
}

export async function releaseCredits(userId: string, count: number) {
  if (count <= 0) return;
  await refundUserCredit(userId, 0, count);
}

export async function releaseReservation(
  userId: string,
  billingSource: "QUOTA" | "CREDIT",
  count = 1,
) {
  if (billingSource === "CREDIT") {
    await releaseCredits(userId, count);
  } else {
    await releaseQuota(userId, count);
  }
}

export async function releaseReservationSplit(userId: string, split: ReservationSplit) {
  if (split.fromQuota > 0) await releaseQuota(userId, split.fromQuota);
  if (split.fromCredits > 0) await releaseCredits(userId, split.fromCredits);
}

/** @deprecated Prefer reserveQuota at create */
export async function incrementQuota(userId: string, count: number) {
  await reserveQuota(userId, count);
}

export function getInitialQuotaResetAt(): Date {
  return getNextMonthlyQuotaReset();
}