/** Token bucket: `perSecond` refill, up to `burst` stored. One per socket. */
export class TokenBucket {
  private tokens: number;
  private last: number;

  constructor(
    private readonly perSecond: number,
    private readonly burst: number,
    private readonly now: () => number,
  ) {
    this.tokens = burst;
    this.last = now();
  }

  take(cost = 1): boolean {
    const t = this.now();
    const elapsed = Math.max(0, t - this.last) / 1000;
    this.last = t;
    this.tokens = Math.min(this.burst, this.tokens + elapsed * this.perSecond);
    if (this.tokens < cost) return false;
    this.tokens -= cost;
    return true;
  }
}

/** Counts live connections per IP. Bars put every phone behind one Wi-Fi NAT, so the cap is generous. */
export class ConnectionCounter {
  private readonly counts = new Map<string, number>();

  constructor(private readonly max: number) {}

  tryAdd(ip: string): boolean {
    const n = this.counts.get(ip) ?? 0;
    if (n >= this.max) return false;
    this.counts.set(ip, n + 1);
    return true;
  }

  remove(ip: string): void {
    const n = (this.counts.get(ip) ?? 1) - 1;
    if (n <= 0) this.counts.delete(ip);
    else this.counts.set(ip, n);
  }
}

/** Joins create players and database rows, so they get their own, tighter limits. */
export const JOINS_PER_SOCKET = { perSecond: 3 / 60, burst: 4 } as const;
/** A busy bar's whole crowd can join through one Wi-Fi NAT address within a minute or two. */
export const JOINS_PER_IP = { perSecond: 2, burst: 60 } as const;

/** One token bucket per key (e.g. client address), pruned when idle so memory stays bounded. */
export class KeyedBuckets {
  private readonly buckets = new Map<string, { tokens: number; last: number }>();

  constructor(
    private readonly perSecond: number,
    private readonly burst: number,
    private readonly maxKeys = 50_000,
  ) {}

  take(key: string, now: number, cost = 1): boolean {
    let b = this.buckets.get(key);
    if (!b) {
      if (this.buckets.size >= this.maxKeys) this.prune(now, 0);
      b = { tokens: this.burst, last: now };
      this.buckets.set(key, b);
    }
    b.tokens = Math.min(this.burst, b.tokens + (Math.max(0, now - b.last) / 1000) * this.perSecond);
    b.last = now;
    if (b.tokens < cost) return false;
    b.tokens -= cost;
    return true;
  }

  /** Drop buckets idle long enough to have refilled completely (they behave like new ones). */
  prune(now: number, idleMs = (this.burst / this.perSecond) * 1000): void {
    for (const [key, b] of this.buckets) if (now - b.last >= idleMs) this.buckets.delete(key);
  }

  get size(): number {
    return this.buckets.size;
  }
}