import { BaseKVDriver, type BatchOp, type KVDriver, type ScanOptions } from '../src/kv/driver.js'
import { InMemoryKVDriver } from '../src/kv/memory.js'

/** A minimal driver implementing only the four required methods, over a plain Map. */
export class MinimalDriver extends BaseKVDriver {
  readonly map = new Map<string, string>()
  async getMany(keys: string[]) {
    return keys.map((k) => (this.map.has(k) ? JSON.parse(this.map.get(k)!) : undefined))
  }
  async setMany(entries: [string, unknown][]) {
    for (const [k, v] of entries) this.map.set(k, JSON.stringify(v))
  }
  async deleteMany(keys: string[]) {
    for (const k of keys) this.map.delete(k)
  }
  async *scan({ gte, lt, limit = Infinity, reverse, values }: ScanOptions): AsyncGenerator<[string, unknown]> {
    let keys = [...this.map.keys()].filter((k) => (gte === undefined || k >= gte) && (lt === undefined || k < lt)).sort()
    if (reverse) keys.reverse()
    keys = keys.slice(0, limit)
    for (const k of keys) {
      if (!this.map.has(k)) continue
      yield [k, values ? JSON.parse(this.map.get(k)!) : undefined]
    }
  }
}

export type Capability = 'setIfAbsent' | 'batch' | 'deleteRange'

/**
 * Wraps a driver, counting calls per method and optionally hiding optional capabilities.
 * `scan` counts each scan started (i.e. each seek).
 */
export class CountingDriver implements KVDriver {
  calls: Record<string, number> = {}
  setIfAbsent?: KVDriver['setIfAbsent']
  batch?: KVDriver['batch']
  deleteRange?: KVDriver['deleteRange']

  constructor(
    readonly inner: KVDriver = new InMemoryKVDriver(),
    hide: Capability[] = [],
  ) {
    if (inner.setIfAbsent && !hide.includes('setIfAbsent'))
      this.setIfAbsent = (e: [string, unknown][]) => (this.#count('setIfAbsent'), inner.setIfAbsent!(e))
    if (inner.batch && !hide.includes('batch')) this.batch = (ops: BatchOp[]) => (this.#count('batch'), inner.batch!(ops))
    if (inner.deleteRange && !hide.includes('deleteRange'))
      this.deleteRange = (gte: string, lt: string) => (this.#count('deleteRange'), inner.deleteRange!(gte, lt))
  }

  #count(m: string): void {
    this.calls[m] = (this.calls[m] ?? 0) + 1
  }

  reset(): void {
    this.calls = {}
  }

  get total(): number {
    return Object.values(this.calls).reduce((a, b) => a + b, 0)
  }

  getMany(keys: string[]) {
    this.#count('getMany')
    return this.inner.getMany(keys)
  }
  setMany(entries: [string, unknown][]) {
    this.#count('setMany')
    return this.inner.setMany(entries)
  }
  deleteMany(keys: string[]) {
    this.#count('deleteMany')
    return this.inner.deleteMany(keys)
  }
  scan(opts: ScanOptions) {
    this.#count('scan')
    return this.inner.scan(opts)
  }
}

export async function collect<T>(it: AsyncIterable<T>): Promise<T[]> {
  const out: T[] = []
  for await (const x of it) out.push(x)
  return out
}