import { KVStorageHttp } from '../../src/kv-storage/kv-storage-http.js'
import { KVStorageHttpServer, type KVStorageHttpServerOptions } from '../../src/kv-storage/kv-storage-http-server.js'
import { KVStorageInMemory } from '../../src/kv-storage/kv-storage-in-memory.js'
import type { RequestSigner } from '../../src/utils/request-auth.js'
import { KVStorageCounting } from '../helpers.js'

/**
 * An in-process server with a network switch. While `online` is false, new requests fail like a
 * dropped connection (open SSE streams stay up, so pushes can still be simulated).
 * `dropStreams()` ends every open SSE stream, as a proxy or server restart would.
 */
export class TestNetwork {
  online = true
  readonly store = new KVStorageCounting(new KVStorageInMemory())
  server: KVStorageHttpServer
  private readonly clients: KVStorageHttp[] = []
  private readonly streams = new Set<() => void>()

  constructor(options: Partial<KVStorageHttpServerOptions> = {}) {
    this.server = new KVStorageHttpServer(this.store, { auth: 'none', ...options })
  }

  client(signer?: RequestSigner): KVStorageHttp {
    const client = new KVStorageHttp({
      url: 'http://server',
      signer,
      fetch: async (url, init) => {
        if (!this.online) throw new TypeError('fetch failed')
        const response = await this.server.handle(new Request(url, init))
        return url.endsWith('/events') && response.body ? this.droppable(response) : response
      },
    })
    this.clients.push(client)
    return client
  }

  dropStreams(): void {
    for (const drop of this.streams) drop()
    this.streams.clear()
  }

  private droppable(response: Response): Response {
    const reader = response.body!.getReader()
    let drop!: () => void
    const body = new ReadableStream<Uint8Array>({
      start: (controller) => {
        drop = () => {
          controller.close()
          void reader.cancel()
        }
        this.streams.add(drop)
      },
      pull: async (controller) => {
        const { done, value } = await reader.read()
        if (done) controller.close()
        else controller.enqueue(value)
      },
      cancel: async () => {
        this.streams.delete(drop)
        await reader.cancel()
      },
    })
    return new Response(body, { status: response.status, headers: response.headers })
  }

  async close(): Promise<void> {
    await Promise.all(this.clients.map((c) => c.close()))
  }
}