import { describe, it, expect } from 'vitest'
import { createMemoryStore } from '../src/memory-store.js'

const rule = { endpoint: 'llm', by: 'ip', capacity: 2, refillPerSec: 1 }

describe('memory-store', () => {
  it('take 在 capacity 内连续放行、超限拒绝', async () => {
    const store = createMemoryStore()
    const a = await store.take('k', 1, rule, 0)
    const b = await store.take('k', 1, rule, 0)
    const c = await store.take('k', 1, rule, 0)
    expect(a.allowed).toBe(true)
    expect(b.allowed).toBe(true)
    expect(c.allowed).toBe(false)
    expect(c.retryAfterMs).toBeGreaterThan(0)
  })

  it('时间推进后恢复放行', async () => {
    const store = createMemoryStore()
    await store.take('k', 1, rule, 0)
    await store.take('k', 1, rule, 0)
    expect((await store.take('k', 1, rule, 0)).allowed).toBe(false)
    // 1s 后补 1 个
    expect((await store.take('k', 1, rule, 1000)).allowed).toBe(true)
  })

  it('peek 不消耗令牌', async () => {
    const store = createMemoryStore()
    await store.take('k', 1, rule, 0) // 剩 1
    const p1 = await store.peek('k', rule, 0)
    const p2 = await store.peek('k', rule, 0)
    expect(p1.remaining).toBe(1)
    expect(p2.remaining).toBe(1)
    // peek 后仍能再 take 1 次成功
    expect((await store.take('k', 1, rule, 0)).allowed).toBe(true)
  })

  it('并发顺序调用不超发（同步原子，不超过 capacity）', async () => {
    const store = createMemoryStore()
    const results = await Promise.all(
      Array.from({ length: 5 }, () => store.take('k', 1, rule, 0)),
    )
    const allowed = results.filter((r) => r.allowed).length
    expect(allowed).toBe(2) // 恰好 capacity 个
  })

  it('不同 key 互不串桶', async () => {
    const store = createMemoryStore()
    await store.take('a', 1, rule, 0)
    await store.take('a', 1, rule, 0)
    expect((await store.take('a', 1, rule, 0)).allowed).toBe(false)
    expect((await store.take('b', 1, rule, 0)).allowed).toBe(true)
  })
})