src/shared/utils/rate-limiter.ts

import { logger } from "./logger";

interface BucketState {
  tokens: number;
  lastRefill: number;
}

/**
 * Token bucket rate limiter per provider.
 * Controls how many API calls can be made per time window.
 */
export class RateLimiter {
  private buckets = new Map<string, BucketState>();

  constructor(
    private readonly maxTokens: number,
    private readonly refillRateMs: number // how often to fully refill
  ) {}

  /**
   * Returns true if the call is allowed, false if rate limit exceeded.
   */
  tryConsume(provider: string, tokens = 1): boolean {
    const now = Date.now();
    let bucket = this.buckets.get(provider);

    if (!bucket) {
      bucket = { tokens: this.maxTokens, lastRefill: now };
      this.buckets.set(provider, bucket);
    }

    // Refill based on elapsed time
    const elapsed = now - bucket.lastRefill;
    if (elapsed >= this.refillRateMs) {
      bucket.tokens = this.maxTokens;
      bucket.lastRefill = now;
    }

    if (bucket.tokens < tokens) {
      logger.warn({ provider, tokensLeft: bucket.tokens }, `[RateLimit] ${provider} throttled`);
      return false;
    }

    bucket.tokens -= tokens;
    return true;
  }

  /**
   * Wait until a token is available (blocking version).
   */
  async consume(provider: string, tokens = 1): Promise<void> {
    while (!this.tryConsume(provider, tokens)) {
      await new Promise((r) => setTimeout(r, 500));
    }
  }
}

// ─── Daily quota tracker (persisted in memory, resets at midnight) ────────

interface DailyQuota {
  count: number;
  date: string; // YYYY-MM-DD
}

const dailyQuotas = new Map<string, DailyQuota>();

function todayStr(): string {
  return new Date().toISOString().split("T")[0];
}

export function checkDailyQuota(key: string, limit: number): boolean {
  const today = todayStr();
  const quota = dailyQuotas.get(key);

  if (!quota || quota.date !== today) {
    dailyQuotas.set(key, { count: 0, date: today });
    return true;
  }

  if (quota.count >= limit) {
    logger.warn({ key, count: quota.count, limit }, `[DailyQuota] ${key} limit reached`);
    return false;
  }
  return true;
}

export function incrementDailyQuota(key: string): void {
  const today = todayStr();
  const quota = dailyQuotas.get(key) ?? { count: 0, date: today };
  if (quota.date !== today) {
    quota.count = 0;
    quota.date = today;
  }
  quota.count += 1;
  dailyQuotas.set(key, quota);
}

// Pre-configured limiters for each provider
export const serperLimiter = new RateLimiter(10, 60_000);   // 10 req/min
export const hunterLimiter = new RateLimiter(5, 60_000);    // 5 req/min
export const snovLimiter = new RateLimiter(5, 60_000);      // 5 req/min
export const reoonLimiter = new RateLimiter(10, 60_000);    // 10 req/min
export const playwrightLimiter = new RateLimiter(3, 10_000); // 3 pages per 10s