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

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

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

/** Stored in PlanEntitlement.limitValue — feature explicitly disabled via admin. */
export const ENTITLEMENT_UNAVAILABLE = -1;

/** 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,
    // Multi-step pipeline charges ~4–9 units per deck → low monthly budget
    slide_generate: 10,
    excalidraw_generate: 5,
    ai_flashcard: 5,
    voice_transcribe: 20,
    publish_enhance: 2,
  },
  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,
    // ~7 units/deck → ~20 decks/month at 140 units
    slide_generate: 140,
    excalidraw_generate: 40,
    ai_flashcard: 100,
    voice_transcribe: 500,
    publish_enhance: 15,
  },
  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,
    // ~7 units/deck → ~70 decks/month
    slide_generate: 500,
    excalidraw_generate: 200,
    ai_flashcard: 'unlimited',
    voice_transcribe: 'unlimited',
    publish_enhance: 100,
  },
  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',
    publish_enhance: '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 === ENTITLEMENT_UNAVAILABLE
          ? ENTITLEMENT_UNAVAILABLE
          : row.limitValue;
  }

  return map;
}

/** Merge DB entitlements with FALLBACK for features added after admin seeding.
 *  Does NOT override ENTITLEMENT_UNAVAILABLE (-1) or explicit limits. */
function mergeWithFallback(map: LimitMap): LimitMap {
  const fallback = fallbackToLimitMap();
  for (const tier of Object.keys(fallback) as SubscriptionTier[]) {
    for (const [feature, limit] of Object.entries(fallback[tier])) {
      const current = map[tier][feature];
      if (current === undefined) {
        map[tier][feature] = limit;
      }
    }
  }
  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 = mergeWithFallback(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 || limit === ENTITLEMENT_UNAVAILABLE) 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 === ENTITLEMENT_UNAVAILABLE
          ? ('unavailable' as const)
          : 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 };