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';
import {
  getLimitAsync,
  getTierFeaturesAsync,
  invalidateEntitlementCache,
  TIER_LIMITS,
  type SubscriptionTier,
} from './plan-entitlements';

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

export class QuotaServiceUnavailableError extends Error {
  code = 'QUOTA_SERVICE_UNAVAILABLE';

  constructor(message = 'Quota service temporarily unavailable') {
    super(message);
  }
}

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 TTL_SECONDS = 90 * 24 * 60 * 60;

function getPeriodDates(): { period: string; periodStart: Date; periodEnd: Date } {
  const period = getCurrentPeriodKey();
  const periodStart = new Date(`${period}-01T00:00:00.000Z`);
  const periodEnd = new Date(Date.UTC(periodStart.getUTCFullYear(), periodStart.getUTCMonth() + 1, 1));
  return { period, periodStart, periodEnd };
}

async function getDatabaseUsageCounts(
  userId: string,
  features: string[],
): Promise<Record<string, number>> {
  if (features.length === 0) return {};
  const { periodStart } = getPeriodDates();
  const rows = await prisma.usageLog.findMany({
    where: { userId, periodStart, feature: { in: features } },
    select: { feature: true, requestsCount: true },
  });
  return Object.fromEntries(rows.map((r) => [r.feature, r.requestsCount]));
}

async function mirrorUsageCountToDatabase(
  userId: string,
  feature: string,
  requestsCount: number,
): Promise<void> {
  const { periodStart, periodEnd } = getPeriodDates();
  await prisma.usageLog.upsert({
    where: {
      userId_feature_periodStart: { userId, feature, periodStart },
    },
    create: {
      userId,
      feature,
      periodStart,
      periodEnd,
      requestsCount,
      tokensUsed: 0,
      syncedAt: new Date(),
    },
    update: {
      requestsCount,
      syncedAt: new Date(),
    },
  }).catch((err) => {
    console.error('[entitlements] Failed to mirror usage to DB:', err);
  });
}

/** Fallback when Redis is unavailable — atomic increment in PostgreSQL. */
async function reserveUsageInDatabase(
  userId: string,
  feature: string,
  limit: number,
): Promise<number> {
  const { periodStart, periodEnd } = getPeriodDates();

  return prisma.$transaction(async (tx) => {
    const existing = await tx.usageLog.findUnique({
      where: { userId_feature_periodStart: { userId, feature, periodStart } },
    });
    const current = existing?.requestsCount ?? 0;
    if (current >= limit) return -1;

    const updated = await tx.usageLog.upsert({
      where: { userId_feature_periodStart: { userId, feature, periodStart } },
      create: {
        userId,
        feature,
        periodStart,
        periodEnd,
        requestsCount: 1,
        tokensUsed: 0,
        syncedAt: new Date(),
      },
      update: {
        requestsCount: { increment: 1 },
        syncedAt: new Date(),
      },
    });

    return updated.requestsCount;
  });
}

function parseReserveResult(raw: unknown): number {
  const n = typeof raw === 'number' ? raw : Number(raw);
  return Number.isFinite(n) ? n : NaN;
}

function shouldFailClosedOnRedisError(): boolean {
  return process.env.NODE_ENV === 'production';
}

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
`;

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 = await getLimitAsync(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 redisCurrent = parseRedisInt(currentStr);
    const dbCounts = await getDatabaseUsageCounts(userId, [feature]);
    const current = Math.max(redisCurrent, dbCounts[feature] ?? 0);
    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:', err);
    if (shouldFailClosedOnRedisError()) {
      return {
        allowed: false,
        remaining: 0,
        limit,
        tier,
        reason: 'SERVICE_UNAVAILABLE',
        message: 'Quota service temporarily unavailable. Please try again later.',
      };
    }
    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 = await getLimitAsync(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 raw = await redis.eval(
      RESERVE_LUA,
      1,
      key,
      String(limit),
      String(TTL_SECONDS),
    );
    const newCount = parseReserveResult(raw);

    if (newCount === -1) {
      const byokConfigured = await hasAnyActiveByok(userId);
      throw new QuotaExceededError(
        tier === 'BASIC' ? 'PRO' : 'BUSINESS',
        feature,
        limit,
        limit,
        byokConfigured,
        { currentTier: tier },
      );
    }

    if (!Number.isFinite(newCount) || newCount < 1) {
      throw new Error(`Invalid Redis reserve result: ${String(raw)}`);
    }

    await mirrorUsageCountToDatabase(userId, feature, newCount);
  } catch (err) {
    if (err instanceof QuotaExceededError) throw err;

    console.error('[entitlements] Redis reserve failed, trying DB fallback:', err);

    try {
      const dbCount = await reserveUsageInDatabase(userId, feature, limit);
      if (dbCount === -1) {
        const byokConfigured = await hasAnyActiveByok(userId);
        throw new QuotaExceededError(
          tier === 'BASIC' ? 'PRO' : 'BUSINESS',
          feature,
          limit,
          limit,
          byokConfigured,
          { currentTier: tier },
        );
      }
      return;
    } catch (dbErr) {
      if (dbErr instanceof QuotaExceededError) throw dbErr;
      console.error('[entitlements] DB reserve fallback failed:', dbErr);
    }

    if (shouldFailClosedOnRedisError()) {
      throw new QuotaServiceUnavailableError();
    }
  }
}

/** 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 features = await getTierFeaturesAsync(tier);
  const period = getCurrentPeriodKey();

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

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

  try {
    const values = await redis.mget(...keys);
    const dbCounts = await getDatabaseUsageCounts(userId, features);

    const result: Record<string, { remaining: number; limit: number; used: number }> = {};
    for (let i = 0; i < features.length; i++) {
      const feature = features[i];
      const limit = (await getLimitAsync(tier, feature)) ?? 0;
      const redisCurrent = parseRedisInt(values[i]);
      const dbCurrent = dbCounts[feature] ?? 0;
      const current = Math.max(redisCurrent, dbCurrent);
      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 = (await getLimitAsync(tier, feature)) ?? 0;
      result[feature] = { remaining: limit, limit, used: 0 };
    }
    return result;
  }
}

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

/** @deprecated Use getLimitAsync — sync helper for tests only */
export async function getLimit(tier: SubscriptionTier, feature: string): Promise<number | undefined> {
  return getLimitAsync(tier, feature);
}