Momento/memento-note/lib/plan-entitlements.ts

import { prisma } from './prisma';
import { VALID_FEATURES } from './quota-utils';

export type SubscriptionTier = 'BASIC' | 'PRO' | 'BUSINESS' | 'ENTERPRISE';

/** Hardcoded defaults — used when DB is empty or unavailable. */
export const FALLBACK_TIER_LIMITS: Record<
  SubscriptionTier,
  Record<string, number | 'unlimited'>
> = {
  BASIC: {
    semantic_search: 30,
    auto_tag: 15,
    auto_title: 5,
    brainstorm_create: 1,
    brainstorm_expand: 10,
    brainstorm_enrich: 20,
    suggest_charts: 5,
    ai_flashcard: 5,
    voice_transcribe: 20,
  },
  PRO: {
    semantic_search: 200,
    auto_tag: 500,
    auto_title: 200,
    reformulate: 50,
    chat: 50,
    brainstorm_create: 5,
    brainstorm_expand: 100,
    brainstorm_enrich: 200,
    suggest_charts: 50,
    slide_generate: 20,
    excalidraw_generate: 20,
    ai_flashcard: 100,
    voice_transcribe: 500,
  },
  BUSINESS: {
    semantic_search: 1000,
    auto_tag: 1000,
    auto_title: 1000,
    reformulate: 500,
    chat: 500,
    brainstorm_create: 'unlimited',
    brainstorm_expand: 500,
    brainstorm_enrich: 1000,
    suggest_charts: 200,
    slide_generate: 100,
    excalidraw_generate: 100,
    ai_flashcard: 'unlimited',
    voice_transcribe: 'unlimited',
  },
  ENTERPRISE: {
    semantic_search: 'unlimited',
    auto_tag: 'unlimited',
    auto_title: 'unlimited',
    reformulate: 'unlimited',
    chat: 'unlimited',
    brainstorm_create: 'unlimited',
    brainstorm_expand: 'unlimited',
    brainstorm_enrich: 'unlimited',
    suggest_charts: 'unlimited',
    slide_generate: 'unlimited',
    excalidraw_generate: 'unlimited',
    ai_flashcard: 'unlimited',
    voice_transcribe: 'unlimited',
  },
};

const CACHE_TTL_MS = 60_000;

type LimitMap = Record<SubscriptionTier, Record<string, number | undefined>>;

let cachedLimits: LimitMap | null = null;
let cacheExpiresAt = 0;

function fallbackToLimitMap(): LimitMap {
  const map = {} as LimitMap;
  for (const tier of Object.keys(FALLBACK_TIER_LIMITS) as SubscriptionTier[]) {
    map[tier] = {};
    for (const [feature, limit] of Object.entries(FALLBACK_TIER_LIMITS[tier])) {
      map[tier][feature] = limit === 'unlimited' ? Infinity : limit;
    }
  }
  return map;
}

function buildLimitMapFromRows(
  rows: Array<{ tier: SubscriptionTier; feature: string; limitValue: number | null }>,
): LimitMap {
  const map = {} as LimitMap;
  for (const tier of ['BASIC', 'PRO', 'BUSINESS', 'ENTERPRISE'] as SubscriptionTier[]) {
    map[tier] = {};
  }

  for (const row of rows) {
    map[row.tier][row.feature] =
      row.limitValue === null ? Infinity : row.limitValue;
  }

  return map;
}

export function invalidateEntitlementCache(): void {
  cachedLimits = null;
  cacheExpiresAt = 0;
}

async function loadLimitMap(): Promise<LimitMap> {
  const now = Date.now();
  if (cachedLimits && now < cacheExpiresAt) {
    return cachedLimits;
  }

  try {
    const rows = await prisma.planEntitlement.findMany({
      select: { tier: true, feature: true, limitValue: true },
    });

    if (rows.length === 0) {
      cachedLimits = fallbackToLimitMap();
    } else {
      cachedLimits = buildLimitMapFromRows(rows as Array<{
        tier: SubscriptionTier;
        feature: string;
        limitValue: number | null;
      }>);
    }
  } catch (err) {
    console.error('[plan-entitlements] DB load failed, using fallback:', err);
    cachedLimits = fallbackToLimitMap();
  }

  cacheExpiresAt = now + CACHE_TTL_MS;
  return cachedLimits;
}

export async function getLimitAsync(
  tier: SubscriptionTier,
  feature: string,
): Promise<number | undefined> {
  const map = await loadLimitMap();
  const limit = map[tier]?.[feature];
  if (limit === undefined) return undefined;
  if (limit === Infinity) return Infinity;
  return limit;
}

export async function getTierFeaturesAsync(
  tier: SubscriptionTier,
): Promise<string[]> {
  const map = await loadLimitMap();
  const fromDb = Object.keys(map[tier] ?? {});
  if (fromDb.length > 0) return fromDb;
  return Object.keys(FALLBACK_TIER_LIMITS[tier]);
}

export async function getAllEntitlementsForAdmin(): Promise<
  Array<{ tier: SubscriptionTier; feature: string; limitValue: number | null; mode: 'limited' | 'unlimited' | 'unavailable' }>
> {
  const rows = await prisma.planEntitlement.findMany({
    orderBy: [{ tier: 'asc' }, { feature: 'asc' }],
  });

  if (rows.length > 0) {
    return rows.map((r) => ({
      tier: r.tier as SubscriptionTier,
      feature: r.feature,
      limitValue: r.limitValue,
      mode:
        r.limitValue === null
          ? ('unlimited' as const)
          : ('limited' as const),
    }));
  }

  const seeded: Array<{
    tier: SubscriptionTier;
    feature: string;
    limitValue: number | null;
    mode: 'limited' | 'unlimited' | 'unavailable';
  }> = [];

  for (const tier of Object.keys(FALLBACK_TIER_LIMITS) as SubscriptionTier[]) {
    for (const feature of VALID_FEATURES) {
      const raw = FALLBACK_TIER_LIMITS[tier][feature];
      if (raw === undefined) continue;
      seeded.push({
        tier,
        feature,
        limitValue: raw === 'unlimited' ? null : raw,
        mode: raw === 'unlimited' ? 'unlimited' : 'limited',
      });
    }
  }

  return seeded;
}

export { FALLBACK_TIER_LIMITS as TIER_LIMITS };