import {type FileHandle, open, unlink} from 'node:fs/promises'

import baseDebug from '../debug.js'

const debug = baseDebug.extend('bufferThroughFile')

const CHUNK_SIZE = 1024

/**
 * Takes a source stream that will be drained and written to the provided file name as fast as possible.
 * and returns a function that can be called to create multiple readable stream on top of the buffer file.
 * It will start pulling data from the source stream once the first readableStream is created, writing to the buffer file in the background.
 * The readable streams and can be read at any rate (but will not receive data faster than the buffer file is written to).
 * Note: by default, buffering will run to completion, and this may prevent the process from exiting after done reading from the
 * buffered streams. To stop writing to the buffer file, an AbortSignal can be provided and once it's controller aborts, the buffer file will
 * stop and the file will be removed unless keepFile is set. After the signal is aborted, no new buffered readers can be created.
 *
 * @param source - The source readable stream. Will be drained as fast as possible.
 * @param filename - The filename to write to.
 * @param options - Optional AbortSignal to stop writing to the buffer file.
 * @returns A function that can be called multiple times to create a readable stream on top of the buffer file.
 */
export function bufferThroughFile(
  source: globalThis.ReadableStream<Uint8Array>,
  filename: string,
  options?: {keepFile?: boolean; signal: AbortSignal},
) {
  const signal = options?.signal

  let writeHandle: FileHandle
  let readHandle: Promise<FileHandle> | null

  // Whether the all data has been written to the buffer file.
  let bufferDone = false

  let removal: Promise<void> | undefined
  function removeBufferFile() {
    if (options?.keepFile === true) return
    removal ??= unlink(filename)
    return removal
  }

  let abortCleanup: Promise<void> | undefined
  signal?.addEventListener(
    'abort',
    () => {
      debug('Aborting bufferThroughFile')
      abortCleanup = (async () => {
        // An abort can arrive while open() is still pending. Wait for ownership of the file
        // before closing it, and prevent start() from opening a read handle afterwards.
        await ready
        if (!writeHandle) return
        await Promise.all([
          writeHandle.close(),
          readHandle && readHandle.then((handle) => handle.close()),
        ])
        await removeBufferFile()
      })().catch((error) => {
        debug('Error cleaning up buffer on abort', error)
      })
    },
    {once: true},
  )

  // Number of active readers. When this reaches 0, the read handle will be closed.
  let readerCount = 0
  let ready: Promise<void>

  async function pump(reader: globalThis.ReadableStreamDefaultReader<Uint8Array>) {
    try {
      while (true) {
        const {done, value} = await reader.read()
        if (done || signal?.aborted) {
          // if we're done reading, or the primary reader has been cancelled, stop writing to the buffer file
          return
        }
        await writeHandle.write(value)
      }
    } finally {
      await writeHandle.close()
      bufferDone = true
      reader.releaseLock()
    }
  }

  function createBufferedReader() {
    let totalBytesRead = 0

    return async function tryReadFromBuffer(
      handle: FileHandle,
    ): Promise<{buffer: Uint8Array; bytesRead: number}> {
      let buffer: Uint8Array
      let bytesRead: number
      try {
        ;({buffer, bytesRead} = await handle.read(
          new Uint8Array(CHUNK_SIZE),
          0,
          CHUNK_SIZE,
          totalBytesRead,
        ))
      } catch (error) {
        if (signal?.aborted) {
          // Aborting closes the handles, so a read that was in flight (or started) at that point
          // fails with EBADF. That's an expected consequence of aborting, not a read error.
          debug('Read failed after abort, treating as end of stream')
          return {buffer: new Uint8Array(0), bytesRead: 0}
        }
        throw error
      }
      if (bytesRead === 0 && !bufferDone && !signal?.aborted) {
        debug('Not enough data in buffer file, waiting for more data to be written')
        // we're waiting for more data to be written to the buffer file, try again
        return tryReadFromBuffer(handle)
      }
      totalBytesRead += bytesRead
      return {buffer, bytesRead}
    }
  }

  function init(): Promise<void> {
    if (ready === undefined) {
      ready = (async () => {
        debug('Initializing bufferThroughFile')
        writeHandle = await open(filename, 'w')
        // start pumping data from the source stream to the buffer file
        debug('Start buffering source stream to file')
        // note, don't await this, as it will block the ReadableStream.start() method
        pump(source.getReader())
          .then(() => {
            debug('Buffering source stream to buffer file')
          })
          .catch((error) => {
            debug('Error pumping source stream', error)
          })
      })()
    }
    return ready
  }

  function getReadHandle(): Promise<FileHandle> {
    if (!readHandle) {
      debug('Opening read handle on %s', filename)
      readHandle = open(filename, 'r')
    }
    return readHandle
  }

  function onReaderStart() {
    readerCount++
  }
  async function onReaderEnd() {
    readerCount--
    if (signal?.aborted) {
      await abortCleanup
      return
    }
    if (readerCount === 0 && readHandle) {
      const handle = readHandle
      readHandle = null
      debug('Closing read handle on %s', filename)
      await (await handle).close()
      if (options?.keepFile !== true) {
        debug('Removing buffer file', filename)
        await removeBufferFile()
      }
    }
  }

  return () => {
    const readChunk = createBufferedReader()

    let didEnd = false
    async function onEnd() {
      if (didEnd) {
        return
      }
      didEnd = true
      await onReaderEnd()
    }
    return new ReadableStream<Uint8Array>({
      async cancel() {
        await onEnd()
      },
      async pull(controller) {
        if (signal?.aborted) {
          await onEnd()
          controller.close()
          return
        }
        if (!readHandle) {
          throw new Error('Cannot read from closed handle')
        }
        const {buffer, bytesRead} = await readChunk(await readHandle)
        if (bytesRead === 0 && (bufferDone || signal?.aborted)) {
          debug('Reader done reading from file handle')
          await onEnd()
          controller.close()
        } else {
          controller.enqueue(buffer.subarray(0, bytesRead))
        }
      },
      async start() {
        if (signal?.aborted) {
          throw new Error('Cannot create new buffered readers on aborted stream')
        }
        debug('Reader started reading from file handle')
        onReaderStart()
        await init()
        if (!signal?.aborted) {
          await getReadHandle()
        }
      },
    })
  }
}