import { PassThrough } from 'node:stream'
import ReactDOMServer from 'react-dom/server'
import { isbot } from 'isbot'
import {
  transformPipeableStreamWithRouter,
  transformReadableStreamWithRouter,
} from '@tanstack/router-core/ssr/server'
import type { AnyRouter } from '@tanstack/router-core'
import type { ReadableStream } from 'node:stream/web'
import type { ReactNode } from 'react'

export const renderRouterToStream = async ({
  request,
  router,
  responseHeaders,
  children,
}: {
  request: Request
  router: AnyRouter
  responseHeaders: Headers
  children: ReactNode
}) => {
  if (typeof ReactDOMServer.renderToReadableStream === 'function') {
    const stream = await ReactDOMServer.renderToReadableStream(children, {
      signal: request.signal,
      nonce: router.options.ssr?.nonce,
      progressiveChunkSize: Number.POSITIVE_INFINITY,
    })

    if (isbot(request.headers.get('User-Agent'))) {
      await stream.allReady
    }

    const responseStream = transformReadableStreamWithRouter(
      router,
      stream as unknown as ReadableStream,
    )
    return new Response(responseStream as any, {
      status: router.stores.statusCode.get(),
      headers: responseHeaders,
    })
  }

  if (typeof ReactDOMServer.renderToPipeableStream === 'function') {
    const reactAppPassthrough = new PassThrough()

    try {
      const pipeable = ReactDOMServer.renderToPipeableStream(children, {
        nonce: router.options.ssr?.nonce,
        progressiveChunkSize: Number.POSITIVE_INFINITY,
        ...(isbot(request.headers.get('User-Agent'))
          ? {
              onAllReady() {
                pipeable.pipe(reactAppPassthrough)
              },
            }
          : {
              onShellReady() {
                pipeable.pipe(reactAppPassthrough)
              },
            }),
        onError: (error, info) => {
          console.error('Error in renderToPipeableStream:', error, info)
          // Destroy the passthrough stream on error
          if (!reactAppPassthrough.destroyed) {
            reactAppPassthrough.destroy(
              error instanceof Error ? error : new Error(String(error)),
            )
          }
        },
      })
    } catch (e) {
      console.error('Error in renderToPipeableStream:', e)
      reactAppPassthrough.destroy(e instanceof Error ? e : new Error(String(e)))
    }

    const responseStream = transformPipeableStreamWithRouter(
      router,
      reactAppPassthrough,
    )
    return new Response(responseStream as any, {
      status: router.stores.statusCode.get(),
      headers: responseHeaders,
    })
  }

  throw new Error(
    'No renderToReadableStream or renderToPipeableStream found in react-dom/server. Ensure you are using a version of react-dom that supports streaming.',
  )
}
