import { describe, it, expect, beforeEach } from "vitest";
import { TokenBucket, ConnectLimiter } from "../src/rate-limit.js";

/** A deterministic fake clock. */
function fakeClock(initial = 0) {
  let t = initial;
  return {
    now: () => t,
    advance: (ms: number) => {
      t += ms;
    },
  };
}

describe("TokenBucket (per-socket sustained rate, R4)", () => {
  let clock: ReturnType<typeof fakeClock>;
  let bucket: TokenBucket;

  beforeEach(() => {
    clock = fakeClock();
    bucket = new TokenBucket({ capacity: 5, rate: 5, now: clock.now });
  });

  it("serves the full burst capacity without consuming below zero", () => {
    for (let i = 0; i < 5; i++) {
      expect(bucket.consume("a").ok).toBe(true);
    }
  });

  it("rejects the request past capacity with a positive retryAfterMs", () => {
    for (let i = 0; i < 5; i++) bucket.consume("a");
    const denied = bucket.consume("a");
    expect(denied.ok).toBe(false);
    if (!denied.ok) expect(denied.retryAfterMs).toBeGreaterThan(0);
  });

  it("tracks keys independently", () => {
    for (let i = 0; i < 5; i++) bucket.consume("a");
    expect(bucket.consume("a").ok).toBe(false);
    expect(bucket.consume("b").ok).toBe(true);
  });

  it("refills continuously — the bucket is usable again after (capacity/rate) seconds", () => {
    for (let i = 0; i < 5; i++) bucket.consume("a");
    expect(bucket.consume("a").ok).toBe(false);
    // 5 capacity / 5 per s → a full refill after 1000 ms.
    clock.advance(1000);
    expect(bucket.consume("a").ok).toBe(true);
  });

  it("refills are capped at the capacity (no unbounded accumulation)", () => {
    clock.advance(60_000); // an hour of idleness
    expect(bucket.available("a")).toBe(5); // exactly the capacity, not 3000
  });

  it("drop removes a bucket and size reflects live buckets", () => {
    bucket.consume("a");
    bucket.consume("b");
    expect(bucket.size()).toBe(2);
    bucket.drop("a");
    expect(bucket.size()).toBe(1);
  });

  it("sweep drops buckets idle longer than staleMs, keeps active ones", () => {
    bucket.consume("idle");
    clock.advance(10_000);
    bucket.consume("fresh");
    expect(bucket.sweep(5000)).toBe(1); // only "idle" is stale
    expect(bucket.available("fresh")).toBeGreaterThan(0);
  });
});

describe("ConnectLimiter (connect-time admission, R4)", () => {
  let clock: ReturnType<typeof fakeClock>;
  let limiter: ConnectLimiter;

  beforeEach(() => {
    clock = fakeClock();
    limiter = new ConnectLimiter({ windowMs: 60_000, max: 3, now: clock.now });
  });

  it("admits up to the max per window, then rejects", () => {
    expect(limiter.tryAdmit("1.2.3.4").ok).toBe(true);
    expect(limiter.tryAdmit("1.2.3.4").ok).toBe(true);
    expect(limiter.tryAdmit("1.2.3.4").ok).toBe(true);
    expect(limiter.tryAdmit("1.2.3.4")).toEqual({ ok: false });
  });

  it("tracks keys (IPs) independently", () => {
    limiter.tryAdmit("a");
    limiter.tryAdmit("a");
    limiter.tryAdmit("a");
    expect(limiter.tryAdmit("a").ok).toBe(false);
    expect(limiter.tryAdmit("b").ok).toBe(true);
  });

  it("admits again once the sliding window has passed", () => {
    limiter.tryAdmit("a");
    limiter.tryAdmit("a");
    limiter.tryAdmit("a");
    expect(limiter.tryAdmit("a").ok).toBe(false);
    clock.advance(60_001);
    expect(limiter.tryAdmit("a").ok).toBe(true);
  });

  it("count reflects the in-window hits", () => {
    limiter.tryAdmit("a");
    limiter.tryAdmit("a");
    expect(limiter.count("a")).toBe(2);
  });

  it("sweep drops idle keys", () => {
    limiter.tryAdmit("a");
    clock.advance(120_000);
    expect(limiter.sweep(60_000)).toBe(1);
    expect(limiter.size()).toBe(0);
  });
});