import type {
  RateLimiter,
  RateLimitResult,
  RateLimitConfig,
  RateLimiterState,
} from "../types.js";

/**
 * Token Bucket Rate Limiter
 *
 * Maintains a bucket of tokens that refills at a constant rate.
 * Allows bursts up to bucket capacity while maintaining average rate.
 * Good for APIs where occasional bursts are acceptable.
 */
export class TokenBucketRateLimiter implements RateLimiter {
  private tokens: number;
  private lastRefill: number;
  private config: RateLimitConfig;

  constructor(config: RateLimitConfig) {
    this.config = config;
    this.tokens = config.bucketCapacity;
    this.lastRefill = Date.now();
  }

  private refill(): void {
    const now = Date.now();
    const elapsed = (now - this.lastRefill) / 1000; // seconds
    const tokensToAdd = elapsed * this.config.refillRate;

    this.tokens = Math.min(
      this.config.bucketCapacity,
      this.tokens + tokensToAdd
    );
    this.lastRefill = now;
  }

  check(): RateLimitResult {
    if (!this.config.enabled) {
      return { allowed: true, headers: {} };
    }

    this.refill();

    const headers: Record<string, string> = {
      "x-ratelimit-limit-requests": String(this.config.bucketCapacity),
      "x-ratelimit-remaining-requests": String(Math.floor(this.tokens)),
    };

    if (this.tokens >= 1) {
      this.tokens -= 1;
      return { allowed: true, headers };
    }

    // Calculate time until next token
    const timeToNextTokenSeconds = (1 - this.tokens) / this.config.refillRate;
    const timeToNextTokenMs = Math.ceil(timeToNextTokenSeconds * 1000);

    return {
      allowed: false,
      retryAfter: Math.ceil(timeToNextTokenSeconds),
      headers: {
        ...headers,
        "retry-after": String(Math.ceil(timeToNextTokenSeconds)),
        "retry-after-ms": String(timeToNextTokenMs),
      },
    };
  }

  reset(): void {
    this.tokens = this.config.bucketCapacity;
    this.lastRefill = Date.now();
  }

  getState(): RateLimiterState {
    this.refill(); // Update token count before reporting
    return {
      strategy: "token-bucket",
      tokens: Math.floor(this.tokens),
      capacity: this.config.bucketCapacity,
    };
  }

  updateConfig(config: Partial<RateLimitConfig>): void {
    const oldCapacity = this.config.bucketCapacity;
    this.config = { ...this.config, ...config };

    // If capacity increased, allow immediate use of new tokens
    if (
      config.bucketCapacity !== undefined &&
      config.bucketCapacity > oldCapacity
    ) {
      this.tokens = Math.min(
        config.bucketCapacity,
        this.tokens + (config.bucketCapacity - oldCapacity)
      );
    }
  }
}
