import type { FileHandle } from "node:fs/promises";

type ByteSource = ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>;

const READ_CHUNK_BYTES = 256 * 1024;

export async function* byteChunks(source: ByteSource): AsyncIterable<Uint8Array> {
  if (!("getReader" in source)) {
    yield* source;
    return;
  }
  const reader = source.getReader();
  try {
    for (;;) {
      const { done, value } = await reader.read();
      if (done) return;
      yield value;
    }
  } finally {
    // Stop the producer (e.g. an upload body) when the consumer gives up early, without waiting
    // on a producer that may never confirm.
    void reader.cancel().catch(() => undefined);
    reader.releaseLock();
  }
}

/** Releases a stream nobody will read, so its producer (e.g. an open file) is closed. */
export function discardByteStream(source: ByteSource): void {
  if ("cancel" in source) void source.cancel().catch(() => undefined);
}

/** Collects a stream into one buffer, for destinations that only accept bytes. */
export async function collectBytes(source: ByteSource): Promise<Uint8Array> {
  const chunks: Uint8Array[] = [];
  let size = 0;
  for await (const chunk of byteChunks(source)) {
    chunks.push(chunk);
    size += chunk.byteLength;
  }
  const bytes = new Uint8Array(size);
  let offset = 0;
  for (const chunk of chunks) {
    bytes.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return bytes;
}

export async function writeFully(handle: FileHandle, chunk: Uint8Array) {
  let offset = 0;
  while (offset < chunk.byteLength) {
    const { bytesWritten } = await handle.write(chunk, offset, chunk.byteLength - offset);
    offset += bytesWritten;
  }
}

/** Streams [start, end] of an open file and closes it when done, failed or cancelled. */
export function fileByteStream(handle: FileHandle, start: number, end: number) {
  let position = start;
  let closed = false;
  const close = async () => {
    if (closed) return;
    closed = true;
    await handle.close().catch(() => undefined);
  };
  return new ReadableStream<Uint8Array>({
    async pull(controller) {
      try {
        const length = Math.min(end - position + 1, READ_CHUNK_BYTES);
        if (length <= 0) {
          await close();
          controller.close();
          return;
        }
        const buffer = new Uint8Array(length);
        const { bytesRead } = await handle.read(buffer, 0, length, position);
        if (bytesRead === 0) throw new Error("File ended before the requested range");
        position += bytesRead;
        controller.enqueue(bytesRead === length ? buffer : buffer.subarray(0, bytesRead));
        if (position > end) {
          await close();
          controller.close();
        }
      } catch (error) {
        await close();
        controller.error(error);
      }
    },
    cancel: close,
  });
}