songs2yt/lib/quota.ts

import { Plan } from "@prisma/client";
import { prisma } from "./db";
import { isSelfHostedEdition } from "./edition";
import { getEffectiveQuotaLimit } from "./quota-extensions";
import { getNextMonthlyQuotaReset, getPlanLimits } from "./plans";

const TWELVE_HOURS_MS = 12 * 60 * 60 * 1000;
const CREDIT_BALANCE_MAX = 1000;

function getNextFreeQuotaReset(from: Date = new Date()): Date {
  return new Date(from.getTime() + TWELVE_HOURS_MS);
}

function getNextQuotaResetForPlan(plan: Plan, from: Date = new Date()): Date {
  return plan === "FREE" ? getNextFreeQuotaReset(from) : getNextMonthlyQuotaReset(from);
}

function isMonthlyResetDate(date: Date): boolean {
  return (
    date.getDate() === 1 &&
    date.getHours() === 0 &&
    date.getMinutes() === 0 &&
    date.getSeconds() === 0 &&
    date.getMilliseconds() === 0
  );
}

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;
  if (plan === "PREMIUM") {
    return target.toLocaleDateString(undefined, {
      month: "long",
      day: "numeric",
      year: "numeric",
    });
  }
  return formatQuotaResetCountdown(target);
}

function getQuotaExceededMessage(plan: Plan, resetsAt: Date, videoCredits: number): string {
  if (plan === "PREMIUM") {
    const resetLabel = formatQuotaResetDisplay(plan, resetsAt);
    return `Quota exceeded. Your monthly Pro quota resets on ${resetLabel}. Request an extension in Settings if needed.`;
  }

  const countdown = formatQuotaResetCountdown(resetsAt);
  const creditHint = videoCredits > 0 ? "" : " Credits are not available in this build.";
  return `Not enough free quota or credits. Your free quota resets in ${countdown}.${creditHint}`;
}

export async function ensureQuotaReset(userId: string) {
  const user = await prisma.user.findUniqueOrThrow({ where: { id: userId } });
  const now = new Date();

  if (user.plan === "PREMIUM") {
    if (now >= user.quotaResetAt) {
      return prisma.user.update({
        where: { id: userId },
        data: {
          videosUsed: 0,
          bonusQuota: 0,
          quotaResetAt: getNextMonthlyQuotaReset(now),
        },
      });
    }

    if (!isMonthlyResetDate(user.quotaResetAt)) {
      return prisma.user.update({
        where: { id: userId },
        data: {
          quotaResetAt: getNextMonthlyQuotaReset(now),
        },
      });
    }

    return user;
  }

  if (now >= user.quotaResetAt) {
    return prisma.user.update({
      where: { id: userId },
      data: {
        videosUsed: 0,
        quotaResetAt: getNextQuotaResetForPlan(user.plan, now),
      },
    });
  }

  return user;
}

export async function getQuotaInfo(userId: string) {
  const user = await ensureQuotaReset(userId);
  const limits = getPlanLimits(user.plan);
  const limit = isSelfHostedEdition()
    ? limits.monthlyQuota
    : getEffectiveQuotaLimit(user.plan, user.bonusQuota);
  const remaining = Math.max(0, limit - user.videosUsed);
  const videoCredits =
    isSelfHostedEdition() || user.plan !== "FREE" ? 0 : user.videoCredits;
  return {
    used: user.videosUsed,
    limit,
    baseLimit: limits.monthlyQuota,
    bonusQuota: isSelfHostedEdition() ? 0 : user.bonusQuota,
    remaining,
    videoCredits,
    creditBalanceMax: CREDIT_BALANCE_MAX,
    totalAvailable: remaining + videoCredits,
    resetsAt: user.quotaResetAt.toISOString(),
    resetsIn: formatQuotaResetDisplay(user.plan, user.quotaResetAt),
    plan: user.plan,
    maxBatchSize: limits.maxBatchSize,
    watermarkOptional: limits.watermarkOptional,
    maxResolutionHeight: limits.maxResolutionHeight,
    selfHosted: isSelfHostedEdition(),
  };
}

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

/**
 * Free plan: use included quota first, then pay-as-you-go credits.
 * Pro plan: subscription quota only (PAYG is a separate product).
 */
export function planReservation(
  plan: Plan,
  remainingQuota: number,
  videoCredits: number,
  count: number,
): ReservationSplit | null {
  if (count <= 0) return { fromQuota: 0, fromCredits: 0 };

  if (plan === "PREMIUM") {
    if (count > remainingQuota) return null;
    return { fromQuota: count, fromCredits: 0 };
  }

  const fromQuota = Math.min(count, Math.max(0, remainingQuota));
  const fromCredits = count - fromQuota;
  if (fromCredits > videoCredits) 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,
        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 = getEffectiveQuotaLimit(user.plan, user.bonusQuota);
  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.videoCredits,
      resetsAt: user.quotaResetAt.toISOString(),
    };
  }

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

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

/**
 * Atomically reserve plan quota first, then prepaid credits.
 * Returns how many of each were taken (for per-item billingSource).
 */
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);

  return prisma.$transaction(async (tx) => {
    const rows = await tx.$queryRaw<
      Array<{
        id: string;
        plan: Plan;
        videosUsed: number;
        bonusQuota: number;
        videoCredits: number;
        quotaResetAt: Date;
      }>
    >`SELECT id, plan, "videosUsed", "bonusQuota", "videoCredits", "quotaResetAt" FROM "User" WHERE id = ${userId} FOR UPDATE`;

    const user = rows[0];
    if (!user) throw new Error("User not found");

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

    const limit = getEffectiveQuotaLimit(user.plan, user.bonusQuota);
    const remaining = Math.max(0, limit - user.videosUsed);
    const split = planReservation(user.plan, remaining, user.videoCredits, count);
    if (!split) {
      throw new Error(getQuotaExceededMessage(user.plan, user.quotaResetAt, user.videoCredits));
    }

    await tx.user.update({
      where: { id: userId },
      data: {
        ...(split.fromQuota > 0 ? { videosUsed: { increment: split.fromQuota } } : {}),
        ...(split.fromCredits > 0 ? { videoCredits: { decrement: split.fromCredits } } : {}),
      },
    });

    return split;
  });
}

export async function releaseQuota(userId: string, count: number) {
  if (count <= 0) return;
  await ensureQuotaReset(userId);
  await prisma.$executeRaw`
    UPDATE "User"
    SET "videosUsed" = GREATEST(0, "videosUsed" - ${count})
    WHERE id = ${userId}
  `;
}

export async function releaseCredits(userId: string, count: number) {
  if (count <= 0) return;
  await prisma.$executeRaw`
    UPDATE "User"
    SET "videoCredits" = LEAST(${CREDIT_BALANCE_MAX}, "videoCredits" + ${count})
    WHERE id = ${userId}
  `;
}

/** Release one reserved unit based on how the job item was billed. */
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; kept for callers that still increment on success. */
export async function incrementQuota(userId: string, count: number) {
  await reserveQuota(userId, count);
}

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