// shared capped reader for web streams: 'truncate' stops reading once the
// cap is reached (the last chunk may overshoot slightly), 'reject' throws
// StreamOverflowError the moment the cap is exceeded

export class StreamOverflowError extends Error {}

const concatChunks = (
  chunks: Uint8Array[],
  total: number,
): Uint8Array => {
  const out = new Uint8Array(total)
  let offset = 0
  for (const chunk of chunks) {
    out.set(chunk, offset)
    offset += chunk.byteLength
  }
  return out
}

export const readStreamCapped = async (
  reader: ReadableStreamDefaultReader<Uint8Array> | undefined,
  max: number,
  overflow: 'truncate' | 'reject',
): Promise<Uint8Array> => {
  if (!reader) return new Uint8Array()
  const chunks: Uint8Array[] = []
  let total = 0
  while (overflow === 'reject' || total < max) {
    const {done, value} = await reader.read()
    if (done) break
    total += value.byteLength
    if (overflow === 'reject' && total > max) {
      await reader.cancel().catch(() => {})
      throw new StreamOverflowError()
    }
    chunks.push(value)
  }
  await reader.cancel().catch(() => {})
  return concatChunks(chunks, total)
}
