import { describe, it, expect } from 'vitest'
import {
  initialState,
  refill,
  waitMsFor,
  take,
  peek,
} from '../src/token-bucket.js'

const rule = { capacity: 3, refillPerSec: 1 } // 1 token/s, 桶容量 3

describe('token-bucket', () => {
  it('初始状态是满桶', () => {
    const s = initialState(rule, 1000)
    expect(s.tokens).toBe(3)
    expect(s.lastRefillMs).toBe(1000)
  })

  it('突发：满桶内连续扣减全放行，扣到空', () => {
    let s = initialState(rule, 0)
    for (let i = 0; i < 3; i++) {
      const r = take(s, 1, rule, 0)
      expect(r.allowed).toBe(true)
      s = r.next
    }
    expect(s.tokens).toBe(0)
  })

  it('空桶不放行，retryAfterMs 准确（缺 1 个、速率 1/s → 1000ms）', () => {
    const empty = { tokens: 0, lastRefillMs: 0 }
    const r = take(empty, 1, rule, 0)
    expect(r.allowed).toBe(false)
    expect(r.remaining).toBe(0)
    expect(r.retryAfterMs).toBe(1000)
    expect(r.next.tokens).toBe(0) // 拒绝不扣
  })

  it('按注入 nowMs 推进补充令牌', () => {
    const empty = { tokens: 0, lastRefillMs: 0 }
    // 2.5s 后应补 2.5 个令牌
    const filled = refill(empty, rule, 2500)
    expect(filled.tokens).toBeCloseTo(2.5)
    // 取 1 个应放行
    const r = take(empty, 1, rule, 2500)
    expect(r.allowed).toBe(true)
  })

  it('补充不超过 capacity', () => {
    const empty = { tokens: 0, lastRefillMs: 0 }
    const filled = refill(empty, rule, 10_000) // 想补 10 个
    expect(filled.tokens).toBe(3) // 封顶 capacity
  })

  it('突发到 capacity 后稳态退回 refill 速率', () => {
    let s = initialState(rule, 0)
    // 突发吃光 3 个
    for (let i = 0; i < 3; i++) s = take(s, 1, rule, 0).next
    // 此后每秒只补 1 个：t=1000 放行 1 个，t=1001 再取被拒
    const a = take(s, 1, rule, 1000)
    expect(a.allowed).toBe(true)
    const b = take(a.next, 1, rule, 1001)
    expect(b.allowed).toBe(false)
    expect(b.retryAfterMs).toBeGreaterThan(0)
  })

  it('refillPerSec=0：固定桶耗尽不恢复，等待为无穷', () => {
    const fixed = { capacity: 2, refillPerSec: 0 }
    const empty = { tokens: 0, lastRefillMs: 0 }
    const r = take(empty, 1, fixed, 999_999)
    expect(r.allowed).toBe(false)
    expect(r.retryAfterMs).toBe(Number.POSITIVE_INFINITY)
  })

  it('时钟回拨不倒扣令牌', () => {
    const s = { tokens: 2, lastRefillMs: 5000 }
    const filled = refill(s, rule, 1000) // nowMs 早于 lastRefillMs
    expect(filled.tokens).toBe(2)
  })

  it('peek 不修改余额（无 next 写回语义）', () => {
    const s = { tokens: 1, lastRefillMs: 0 }
    const r = peek(s, 1, rule, 0)
    expect(r.allowed).toBe(true)
    expect(r.remaining).toBe(1)
  })

  it('waitMsFor: deficit<=0 → 0; 速率>0 向上取整', () => {
    expect(waitMsFor(0, rule)).toBe(0)
    expect(waitMsFor(-1, rule)).toBe(0)
    expect(waitMsFor(0.5, { refillPerSec: 1 })).toBe(500)
    expect(waitMsFor(0.001, { refillPerSec: 1 })).toBe(1) // ceil
    expect(waitMsFor(1, { refillPerSec: 0 })).toBe(Number.POSITIVE_INFINITY)
  })
})