Momento/memento-note/lib/entitlements.ts

import { redis } from './redis';
import { prisma } from './prisma';
import { hasAnyActiveByok } from './byok';
import {
  getCurrentPeriodKey,
  getRedisKey,
  parseRedisInt,
  isValidFeature,
  VALID_FEATURES,
} 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,
  amount: number = 1,
): Promise<number> {
  const { periodStart, periodEnd } = getPeriodDates();
  const units = Math.max(1, Math.floor(amount));

  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 + units > limit) return -1;

    const updated = await tx.usageLog.upsert({
      where: { userId_feature_periodStart: { userId, feature, periodStart } },
      create: {
        userId,
        feature,
        periodStart,
        periodEnd,
        requestsCount: units,
        tokensUsed: 0,
        syncedAt: new Date(),
      },
      update: {
        requestsCount: { increment: units },
        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 RELEASE_LUA = `
local amount = tonumber(ARGV[1]) or 1
if amount < 1 then amount = 1 end
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if current <= 0 then return 0 end
local dec = math.min(amount, current)
redis.call('DECRBY', KEYS[1], dec)
return tonumber(redis.call('GET', KEYS[1]) or '0')
`;

/** Libère des crédits précédemment réservés (modèle A). */
export async function releaseUsage(
  userId: string,
  feature: string,
  amount: number = 1,
): Promise<void> {
  if (!isValidFeature(feature)) return;
  const units = Math.max(1, Math.floor(Number(amount) || 1));
  const { releaseCredits } = await import('@/lib/credits');
  await releaseCredits(userId, units, { feature });
}

const RESERVE_LUA = `
local limit = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
local amount = tonumber(ARGV[3]) or 1
if amount < 1 then amount = 1 end
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if current + amount > limit then
  return -1
end
redis.call('INCRBY', KEYS[1], amount)
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 },
    );
  }
}

/** Mappe une feature de quota vers la voie IA (pour le BYOK). */
function featureToAiLane(feature: string): 'chat' | 'tags' | 'embedding' {
  if (feature === 'auto_tag' || feature === 'auto_title') return 'tags';
  if (feature === 'semantic_search') return 'embedding';
  return 'chat';
}

/**
 * Réserve des crédits globaux (modèle A).
 * `amount` = coût en crédits (ex. slides multi-étapes : 1+N).
 * Si l'utilisateur a une clé perso (BYOK) pour la voie concernée, aucun débit.
 * Conserve le nom historique pour ne pas casser tous les appelants.
 */
export async function reserveUsageOrThrow(
  userId: string,
  feature: string,
  amount: number = 1,
): Promise<void> {
  if (!isValidFeature(feature)) {
    throw new QuotaExceededError('PRO', feature, 0, 0, false, { currentTier: 'BASIC' });
  }

  const { getSystemConfig } = await import('@/lib/config');
  const { willUseByokForLane } = await import('@/lib/ai/provider-for-user');
  const config = await getSystemConfig();
  const lane = featureToAiLane(feature);
  const { usedByok } = await willUseByokForLane(lane, config, userId);
  if (usedByok) return;

  const { reserveCreditsOrThrow, resolveCreditCost } = await import('@/lib/credits');
  const cost = resolveCreditCost(feature, { amount });
  await reserveCreditsOrThrow(userId, cost, { feature });
}

/** Units charged for one multi-step slide generation (outline + per-slide expand). */
export function slideGenerateQuotaUnits(slideCount?: number | null): number {
  // 1 outline + N expand calls (clamped). Reflects real multi-call cost.
  const n = typeof slideCount === 'number' && Number.isFinite(slideCount)
    ? Math.min(8, Math.max(3, Math.round(slideCount)))
    : 6;
  return 1 + n;
}

/** 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 type ResetUserUsageOptions = {
  /** Si fourni, ne remet à zéro que ces fonctionnalités. Sinon : toutes. */
  features?: string[];
};

/**
 * Remet à zéro la consommation IA d’un utilisateur pour la période mensuelle en cours
 * (compteurs Redis + lignes UsageLog en base).
 * Réservé aux actions admin.
 */
export async function resetUserUsageForCurrentPeriod(
  userId: string,
  options?: ResetUserUsageOptions,
): Promise<{ period: string; featuresReset: string[]; redisDeleted: number; dbUpdated: number }> {
  if (!userId?.trim()) {
    throw new Error('userId required');
  }

  const period = getCurrentPeriodKey();
  const { periodStart } = getPeriodDates();

  const features = (options?.features?.length
    ? options.features.filter((f) => isValidFeature(f))
    : [...VALID_FEATURES]) as string[];

  if (features.length === 0) {
    return { period, featuresReset: [], redisDeleted: 0, dbUpdated: 0 };
  }

  const keys = features.map((f) => getRedisKey(userId, f));
  let redisDeleted = 0;
  try {
    if (keys.length === 1) {
      redisDeleted = await redis.del(keys[0]);
    } else {
      redisDeleted = await redis.del(...keys);
    }
  } catch (err) {
    console.error('[entitlements] resetUserUsage Redis del failed:', err);
  }

  const dbResult = await prisma.usageLog.updateMany({
    where: {
      userId,
      periodStart,
      feature: { in: features },
    },
    data: {
      requestsCount: 0,
      tokensUsed: 0,
      syncedAt: new Date(),
    },
  });

  return {
    period,
    featuresReset: features,
    redisDeleted,
    dbUpdated: dbResult.count,
  };
}

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);
    if (shouldFailClosedOnRedisError()) {
      throw new QuotaServiceUnavailableError();
    }
    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);
}