Momento/memento-note/lib/entitlements.ts

import { redis } from './redis';
import { prisma } from './prisma';
import { hasAnyActiveByok } from './byok';
import {
  getCurrentPeriodKey,
  getRedisKey,
  parseRedisInt,
  isValidFeature,
} from './quota-utils';

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

export interface EntitlementResult {
  allowed: boolean;
  remaining: number;
  limit: number;
  tier: SubscriptionTier;
  reason?: 'QUOTA_EXCEEDED' | 'TIER_LIMITED' | 'FEATURE_NOT_AVAILABLE';
  message?: string;
  upgradeTier?: 'PRO' | 'BUSINESS';
  byokConfigured?: boolean;
}

export class QuotaExceededError extends Error {
  code = 'QUOTA_EXCEEDED';
  upgradeTier: 'PRO' | 'BUSINESS';
  feature: string;
  currentQuota: number;
  usedQuota: number;
  byokConfigured: boolean;
  billingOwnerId?: string;
  triggeredByUserId?: string;
  isGuestActor?: boolean;
  currentTier?: string;

  constructor(
    upgradeTier: 'PRO' | 'BUSINESS',
    feature: string,
    currentQuota: number,
    usedQuota: number,
    byokConfigured: boolean = false,
    sessionMeta?: {
      billingOwnerId?: string;
      triggeredByUserId?: string;
      isGuestActor?: boolean;
      currentTier?: string;
    },
  ) {
    super(`Quota exceeded for ${feature}`);
    this.upgradeTier = upgradeTier;
    this.feature = feature;
    this.currentQuota = currentQuota;
    this.usedQuota = usedQuota;
    this.byokConfigured = byokConfigured;
    this.billingOwnerId = sessionMeta?.billingOwnerId;
    this.triggeredByUserId = sessionMeta?.triggeredByUserId;
    this.isGuestActor = sessionMeta?.isGuestActor;
    this.currentTier = sessionMeta?.currentTier;
  }

  toJSON() {
    return {
      error: this.code,
      feature: this.feature,
      upgradeTier: this.upgradeTier,
      byokConfigured: this.byokConfigured,
      isGuestActor: this.isGuestActor ?? false,
      currentTier: this.currentTier,
    };
  }
}

const 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,
    slide_generate: 3,
    excalidraw_generate: 3,
  },
  PRO: {
    semantic_search: 100,
    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,
  },
  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,
  },
  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',
  },
};

const TTL_SECONDS = 90 * 24 * 60 * 60;

const INCREMENT_LUA = `
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
local ttl = tonumber(ARGV[1])
redis.call('INCRBY', KEYS[1], 1)
local ttlResult = redis.call('TTL', KEYS[1])
if ttlResult == -1 then
  redis.call('EXPIRE', KEYS[1], ttl)
end
local newCount = tonumber(redis.call('GET', KEYS[1]))
return newCount
`;

const INCREMENT_BY_LUA = `
local count = tonumber(ARGV[1]) or 1
local ttl = tonumber(ARGV[2])
redis.call('INCRBY', KEYS[1], count)
local ttlResult = redis.call('TTL', KEYS[1])
if ttlResult == -1 then
  redis.call('EXPIRE', KEYS[1], ttl)
end
local newCount = tonumber(redis.call('GET', KEYS[1]))
return newCount
`;

const RESERVE_LUA = `
local limit = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if current >= limit then
  return -1
end
redis.call('INCRBY', KEYS[1], 1)
local ttlResult = redis.call('TTL', KEYS[1])
if ttlResult == -1 then
  redis.call('EXPIRE', KEYS[1], ttl)
end
local newCount = tonumber(redis.call('GET', KEYS[1]))
return newCount
`;

function getLimit(tier: SubscriptionTier, feature: string): number | undefined {
  const tierLimits = TIER_LIMITS[tier];
  const limit = tierLimits?.[feature];
  if (limit === 'unlimited') return Infinity;
  if (limit === undefined) return undefined;
  return limit;
}

export async function getUserInfo(
  userId: string,
): Promise<{ tier: SubscriptionTier; status: string; currentPeriodEnd?: Date }> {
  const subscription = await prisma.subscription.findUnique({
    where: { userId },
  });
  if (!subscription) return { tier: 'BASIC', status: 'INACTIVE' };
  return {
    tier: subscription.tier as SubscriptionTier,
    status: subscription.status,
    currentPeriodEnd: subscription.currentPeriodEnd,
  };
}

export async function getEffectiveTier(
  userId: string,
): Promise<SubscriptionTier> {
  const { tier, status, currentPeriodEnd } = await getUserInfo(userId);

  if (status === 'ACTIVE' || status === 'TRIALING') return tier;

  if ((status === 'PAST_DUE' || status === 'CANCELED') && currentPeriodEnd) {
    if (new Date() < new Date(currentPeriodEnd)) return tier;
  }

  return 'BASIC';
}

export async function canUseFeature(
  userId: string,
  feature: string,
): Promise<EntitlementResult> {
  if (!isValidFeature(feature)) {
    return {
      allowed: false,
      remaining: 0,
      limit: 0,
      tier: 'BASIC',
      reason: 'TIER_LIMITED',
      message: `Unknown feature: ${feature}`,
    };
  }

  const tier = await getEffectiveTier(userId);
  const limit = getLimit(tier, feature);

  if (limit === undefined) {
    return {
      allowed: false,
      remaining: 0,
      limit: 0,
      tier,
      reason: 'FEATURE_NOT_AVAILABLE',
      message: `Feature "${feature}" is not available on the ${tier} plan. Upgrade to unlock it.`,
      upgradeTier: tier === 'BASIC' ? 'PRO' : 'BUSINESS',
    };
  }

  if (limit === Infinity) {
    return { allowed: true, remaining: Infinity, limit: Infinity, tier };
  }

  try {
    const key = getRedisKey(userId, feature);
    const currentStr = await redis.get(key);
    const current = parseRedisInt(currentStr);
    const allowed = current < limit;

    if (!allowed) {
      const byokConfigured = await hasAnyActiveByok(userId);
      return {
        allowed: false,
        remaining: 0,
        limit,
        tier,
        reason: 'QUOTA_EXCEEDED',
        upgradeTier: tier === 'BASIC' ? 'PRO' : 'BUSINESS',
        byokConfigured,
      };
    }

    return {
      allowed: true,
      remaining: Math.max(0, limit - current),
      limit,
      tier,
      byokConfigured: await hasAnyActiveByok(userId),
    };
  } catch (err) {
    console.error('[entitlements] Redis unavailable, allowing request (fail-open):', err);
    return { allowed: true, remaining: limit, limit, tier };
  }
}

export async function checkEntitlementOrThrow(
  userId: string,
  feature: string,
): Promise<void> {
  const result = await canUseFeature(userId, feature);

  if (!result.allowed) {
    throw new QuotaExceededError(
      result.upgradeTier ?? (result.tier === 'BASIC' ? 'PRO' : 'BUSINESS'),
      feature,
      result.limit,
      result.limit > 0 ? result.limit - result.remaining : 0,
      result.byokConfigured ?? false,
      { currentTier: result.tier },
    );
  }
}

export async function reserveUsageOrThrow(
  userId: string,
  feature: string,
): Promise<void> {
  if (!isValidFeature(feature)) {
    throw new QuotaExceededError('PRO', feature, 0, 0, false, { currentTier: 'BASIC' });
  }

  const tier = await getEffectiveTier(userId);
  const limit = getLimit(tier, feature);

  if (limit === undefined) {
    throw new QuotaExceededError(
      tier === 'BASIC' ? 'PRO' : 'BUSINESS',
      feature,
      0,
      0,
      false,
      { currentTier: tier },
    );
  }

  if (limit === Infinity) return;

  try {
    const key = getRedisKey(userId, feature);
    const newCount = await redis.eval(
      RESERVE_LUA,
      1,
      key,
      String(limit),
      String(TTL_SECONDS),
    ) as number;

    if (newCount === -1) {
      const byokConfigured = await hasAnyActiveByok(userId);
      throw new QuotaExceededError(
        tier === 'BASIC' ? 'PRO' : 'BUSINESS',
        feature,
        limit,
        limit,
        byokConfigured,
        { currentTier: tier },
      );
    }
  } catch (err) {
    if (err instanceof QuotaExceededError) throw err;
    console.error('[entitlements] Redis unavailable, allowing request (fail-open):', err);
  }
}

/** Host-pays: bill session owner, attach actor metadata for 402 responses (Story 3.4). */
export async function checkSessionEntitlementOrThrow(
  billingOwnerId: string,
  triggeredByUserId: string,
  isGuestActor: boolean,
  feature: string,
): Promise<void> {
  try {
    await reserveUsageOrThrow(billingOwnerId, feature);
  } catch (err) {
    if (err instanceof QuotaExceededError) {
      err.billingOwnerId = billingOwnerId;
      err.triggeredByUserId = triggeredByUserId;
      err.isGuestActor = isGuestActor;
    }
    throw err;
  }
}

export function incrementUsageAsync(userId: string, feature: string, count: number = 1): Promise<void> {
  if (!isValidFeature(feature)) return Promise.resolve();

  const key = getRedisKey(userId, feature);
  return redis.eval(INCREMENT_BY_LUA, 1, key, String(count), String(TTL_SECONDS)).then(() => {}).catch((err) => {
    console.error('[entitlements] Async increment failed:', err);
  });
}

export async function getUserQuotas(
  userId: string,
): Promise<Record<string, { remaining: number; limit: number; used: number }>> {
  const tier = await getEffectiveTier(userId);
  const period = getCurrentPeriodKey();
  const features = Object.keys(TIER_LIMITS[tier]);

  if (features.length === 0) return {};

  const keys = features.map((f) => `usage:${userId}:${f}:${period}`);

  try {
    const values = await redis.mget(...keys);

    const result: Record<string, { remaining: number; limit: number; used: number }> = {};
    for (let i = 0; i < features.length; i++) {
      const feature = features[i];
      const limit = getLimit(tier, feature) ?? 0;
      const current = parseRedisInt(values[i]);
      result[feature] = {
        remaining: limit === Infinity ? Infinity : Math.max(0, limit - current),
        limit,
        used: current,
      };
    }

    return result;
  } catch (err) {
    console.error('[entitlements] getUserQuotas Redis error:', err);
    const result: Record<string, { remaining: number; limit: number; used: number }> = {};
    for (const feature of features) {
      const limit = getLimit(tier, feature) ?? 0;
      result[feature] = { remaining: limit, limit, used: 0 };
    }
    return result;
  }
}

export { TIER_LIMITS, getLimit };
export type { SubscriptionTier };