prmbr-image-mksaas/src/lib/credits.ts

import { randomUUID } from 'crypto';
import { getDb } from '@/db';
import { creditTransaction, userCredit } from '@/db/schema';
import { addDays, isAfter } from 'date-fns';
import { and, asc, eq, or } from 'drizzle-orm';
import {
  CREDIT_EXPIRE_DAYS,
  CREDIT_TRANSACTION_TYPE,
  FREE_MONTHLY_CREDITS,
  REGISTER_GIFT_CREDITS,
} from './constants';

/**
 * Get user's current credit balance
 * @param userId - User ID
 * @returns User's current credit balance
 */
export async function getUserCredits(userId: string): Promise<number> {
  const db = await getDb();
  const record = await db
    .select()
    .from(userCredit)
    .where(eq(userCredit.userId, userId))
    .limit(1);
  return record[0]?.currentCredits || 0;
}

/**
 * Write a credit transaction record
 * @param params - Credit transaction parameters
 */
async function logCreditTransaction({
  userId,
  type,
  amount,
  description,
  paymentId,
  expirationDate,
}: {
  userId: string;
  type: string;
  amount: number;
  description: string;
  paymentId?: string;
  expirationDate?: Date;
}) {
  if (!userId || !type || !description) {
    throw new Error('Invalid params');
  }
  if (!Number.isFinite(amount) || amount === 0) {
    throw new Error('Invalid amount');
  }
  const db = await getDb();
  await db.insert(creditTransaction).values({
    id: randomUUID(),
    userId,
    type,
    amount,
    // remaining amount is the same as amount for earn transactions
    // remaining amount is null for spend transactions
    remainingAmount: amount > 0 ? amount : null,
    description,
    paymentId,
    expirationDate,
    createdAt: new Date(),
    updatedAt: new Date(),
  });
}

/**
 * Add credits (registration, monthly, purchase, etc.)
 * @param params - Credit creation parameters
 */
export async function addCredits({
  userId,
  amount,
  type,
  description,
  paymentId,
  expireDays = CREDIT_EXPIRE_DAYS,
}: {
  userId: string;
  amount: number;
  type: string;
  description: string;
  paymentId?: string;
  expireDays?: number;
}) {
  if (!userId || !type || !description) {
    throw new Error('Invalid params');
  }
  if (!Number.isFinite(amount) || amount <= 0) {
    throw new Error('Invalid amount');
  }
  if (!Number.isFinite(expireDays) || expireDays <= 0) {
    throw new Error('Invalid expire days');
  }
  // Process expired credits first
  await processExpiredCredits(userId);
  // Update user credit balance
  const db = await getDb();
  const current = await db
    .select()
    .from(userCredit)
    .where(eq(userCredit.userId, userId))
    .limit(1);
  const newBalance = (current[0]?.currentCredits || 0) + amount;
  if (current.length > 0) {
    await db
      .update(userCredit)
      .set({
        currentCredits: newBalance,
        lastRefreshAt: new Date(), // TODO: maybe we can not update this field here
        updatedAt: new Date(),
      })
      .where(eq(userCredit.userId, userId));
  } else {
    await db.insert(userCredit).values({
      id: randomUUID(),
      userId,
      currentCredits: newBalance,
      lastRefreshAt: new Date(), // TODO: maybe we can not update this field here
      createdAt: new Date(),
      updatedAt: new Date(),
    });
  }
  // Write credit transaction record
  await logCreditTransaction({
    userId,
    type,
    amount,
    description,
    paymentId,
    // TODO: maybe there is no expiration date for PURCHASE type?
    expirationDate: addDays(new Date(), expireDays),
  });
}

/**
 * Consume credits (FIFO, by expiration)
 * @param params - Credit consumption parameters
 */
export async function consumeCredits({
  userId,
  amount,
  description,
}: {
  userId: string;
  amount: number;
  description: string;
}) {
  if (!userId || !description) {
    throw new Error('Invalid params');
  }
  if (!Number.isFinite(amount) || amount <= 0) {
    throw new Error('Invalid amount');
  }
  // Process expired credits first
  await processExpiredCredits(userId);
  // Check balance
  const balance = await getUserCredits(userId);
  if (balance < amount) {
    console.error(
      `Insufficient credits for user ${userId}, balance: ${balance}, amount: ${amount}, description: ${description}`
    );
    throw new Error('Insufficient credits');
  }
  // FIFO consumption: consume from the earliest unexpired credits first
  const db = await getDb();
  const transactions = await db
    .select()
    .from(creditTransaction)
    .where(
      and(
        eq(creditTransaction.userId, userId),
        or(
          eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.PURCHASE),
          eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.MONTHLY_REFRESH),
          eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.REGISTER_GIFT)
        )
      )
    )
    .orderBy(
      asc(creditTransaction.expirationDate),
      asc(creditTransaction.createdAt)
    );
  // Consume credits
  let left = amount;
  for (const transaction of transactions) {
    if (left <= 0) break;
    const remain = transaction.remainingAmount || 0;
    if (remain <= 0) continue;
    // credits to consume at most in this transaction
    const consume = Math.min(remain, left);
    await db
      .update(creditTransaction)
      .set({
        remainingAmount: remain - consume,
        updatedAt: new Date(),
      })
      .where(eq(creditTransaction.id, transaction.id));
    left -= consume;
  }
  // Update balance
  const current = await db
    .select()
    .from(userCredit)
    .where(eq(userCredit.userId, userId))
    .limit(1);
  const newBalance = (current[0]?.currentCredits || 0) - amount;
  // TODO: there must have one record for this user in userCredit?
  await db
    .update(userCredit)
    .set({ currentCredits: newBalance, updatedAt: new Date() })
    .where(eq(userCredit.userId, userId));
  // Write usage record
  await logCreditTransaction({
    userId,
    type: CREDIT_TRANSACTION_TYPE.USAGE,
    amount: -amount,
    description,
  });
}

/**
 * Process expired credits
 * @param userId - User ID
 */
export async function processExpiredCredits(userId: string) {
  const now = new Date();
  // Get all credit transactions without type EXPIRE
  const db = await getDb();
  const transactions = await db
    .select()
    .from(creditTransaction)
    .where(
      and(
        eq(creditTransaction.userId, userId),
        or(
          // TODO: credits with PURCHASE type can not be expired?
          // eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.PURCHASE),
          eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.MONTHLY_REFRESH),
          eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.REGISTER_GIFT)
        )
      )
    );
  let expiredTotal = 0;
  // Process expired credit transactions
  for (const transaction of transactions) {
    if (
      transaction.expirationDate &&
      isAfter(now, transaction.expirationDate) &&
      !transaction.expirationDateProcessedAt
    ) {
      const remain = transaction.remainingAmount || 0;
      if (remain > 0) {
        expiredTotal += remain;
        await db
          .update(creditTransaction)
          .set({
            remainingAmount: 0,
            expirationDateProcessedAt: now,
            updatedAt: now,
          })
          .where(eq(creditTransaction.id, transaction.id));
      }
    }
  }
  if (expiredTotal > 0) {
    // Deduct expired credits from balance
    const current = await db
      .select()
      .from(userCredit)
      .where(eq(userCredit.userId, userId))
      .limit(1);
    const newBalance = Math.max(
      0,
      (current[0]?.currentCredits || 0) - expiredTotal
    );
    await db
      .update(userCredit)
      .set({ currentCredits: newBalance, updatedAt: now })
      .where(eq(userCredit.userId, userId));
    // Write expire record
    await logCreditTransaction({
      userId,
      type: CREDIT_TRANSACTION_TYPE.EXPIRE,
      amount: -expiredTotal,
      description: `Expire credits: ${expiredTotal}`,
    });
  }
}

/**
 * Add register gift credits
 * @param userId - User ID
 */
export async function addRegisterGiftCredits(userId: string) {
  // Check if user has already received register gift credits
  const db = await getDb();
  const record = await db
    .select()
    .from(creditTransaction)
    .where(
      and(
        eq(creditTransaction.userId, userId),
        eq(creditTransaction.type, CREDIT_TRANSACTION_TYPE.REGISTER_GIFT)
      )
    )
    .limit(1);
  // add register gift credits if user has not received them yet
  if (record.length === 0) {
    await addCredits({
      userId,
      amount: REGISTER_GIFT_CREDITS,
      type: CREDIT_TRANSACTION_TYPE.REGISTER_GIFT,
      description: `Register gift credits: ${REGISTER_GIFT_CREDITS}`,
    });
  }
}

/**
 * Add free monthly credits
 * @param userId - User ID
 */
export async function addMonthlyFreeCredits(userId: string) {
  // Check last refresh time
  const db = await getDb();
  const record = await db
    .select()
    .from(userCredit)
    .where(eq(userCredit.userId, userId))
    .limit(1);
  const now = new Date();
  let canAdd = false;
  // never added credits before
  if (!record[0]?.lastRefreshAt) {
    canAdd = true;
  } else {
    const last = new Date(record[0].lastRefreshAt);
    canAdd =
      now.getMonth() !== last.getMonth() ||
      now.getFullYear() !== last.getFullYear();
  }
  // add credits if it's a new month
  if (canAdd) {
    await addCredits({
      userId,
      amount: FREE_MONTHLY_CREDITS,
      type: CREDIT_TRANSACTION_TYPE.MONTHLY_REFRESH,
      description: `Free monthly credits: ${FREE_MONTHLY_CREDITS} for ${now.getFullYear()}-${now.getMonth() + 1}`,
    });
  }
}