@charlesmatte/granola

/**
 * Token bucket matching Granola's documented limit (burst 25, sustained 5 requests/second).
 * Used by multi-request exports (search) so they stay under the limit instead of leaning on
 * 429 retries. `now`/`sleep` are injectable for tests.
 */
export class TokenBucket {
	private tokens: number
	private last: number
	constructor(
		private readonly capacity = 25,
		private readonly refillPerSecond = 5,
		private readonly now: () => number = () => Date.now(),
		private readonly sleep: (ms: number) => Promise<void> = (ms) => new Promise((r) => setTimeout(r, ms)),
	) {
		this.tokens = capacity
		this.last = now()
	}
	private refill() {
		const t = this.now()
		this.tokens = Math.min(this.capacity, this.tokens + ((t - this.last) / 1000) * this.refillPerSecond)
		this.last = t
	}
	async take(): Promise<void> {
		for (;;) {
			this.refill()
			if (this.tokens >= 1) {
				this.tokens -= 1
				return
			}
			await this.sleep(Math.ceil(((1 - this.tokens) / this.refillPerSecond) * 1000))
		}
	}
}

/** Run `fn` over items with bounded concurrency, preserving input order in the output. */
export async function mapLimit<T, R>(items: Array<T>, concurrency: number, fn: (item: T, index: number) => Promise<R>): Promise<Array<R>> {
	const out = new Array<R>(items.length)
	let next = 0
	const worker = async () => {
		for (;;) {
			const i = next++
			if (i >= items.length) return
			out[i] = await fn(items[i] as T, i)
		}
	}
	await Promise.all(Array.from({ length: Math.min(concurrency, items.length) }, worker))
	return out
}