infiniteQueryBehavior.ts

import type { QueryBehavior } from './query'

import type {
  InfiniteData,
  QueryFunctionContext,
  QueryOptions,
  RefetchQueryFilters,
} from './types'
import { getAbortController } from './utils'

export function infiniteQueryBehavior<
  TQueryFnData,
  TError,
  TData
>(): QueryBehavior<TQueryFnData, TError, InfiniteData<TData>> {
  return {
    onFetch: context => {
      context.fetchFn = () => {
        const refetchPage: RefetchQueryFilters['refetchPage'] | undefined =
          context.fetchOptions?.meta?.refetchPage
        const fetchMore = context.fetchOptions?.meta?.fetchMore
        const pageParam = fetchMore?.pageParam
        const isFetchingNextPage = fetchMore?.direction === 'forward'
        const isFetchingPreviousPage = fetchMore?.direction === 'backward'
        const oldPages = context.state.data?.pages || []
        const oldPageParams = context.state.data?.pageParams || []
        const abortController = getAbortController()
        const abortSignal = abortController?.signal
        let newPageParams = oldPageParams
        let cancelled = false

        // Get query function
        const queryFn =
          context.options.queryFn || (() => Promise.reject('Missing queryFn'))

        const buildNewPages = (
          pages: unknown[],
          param: unknown,
          page: unknown,
          previous?: boolean
        ) => {
          newPageParams = previous
            ? [param, ...newPageParams]
            : [...newPageParams, param]
          return previous ? [page, ...pages] : [...pages, page]
        }

        // Create function to fetch a page
        const fetchPage = (
          pages: unknown[],
          manual?: boolean,
          param?: unknown,
          previous?: boolean
        ): Promise<unknown[]> => {
          if (cancelled) {
            return Promise.reject('Cancelled')
          }

          if (typeof param === 'undefined' && !manual && pages.length) {
            return Promise.resolve(pages)
          }

          const queryFnContext: QueryFunctionContext = {
            queryKey: context.queryKey,
            signal: abortSignal,
            pageParam: param,
            meta: context.meta,
          }

          const queryFnResult = queryFn(queryFnContext)

          const promise = Promise.resolve(queryFnResult).then(page =>
            buildNewPages(pages, param, page, previous)
          )

          return promise
        }

        let promise: Promise<unknown[]>

        // Fetch first page?
        if (!oldPages.length) {
          promise = fetchPage([])
        }

        // Fetch next page?
        else if (isFetchingNextPage) {
          const manual = typeof pageParam !== 'undefined'
          const param = manual
            ? pageParam
            : getNextPageParam(context.options, oldPages)
          promise = fetchPage(oldPages, manual, param)
        }

        // Fetch previous page?
        else if (isFetchingPreviousPage) {
          const manual = typeof pageParam !== 'undefined'
          const param = manual
            ? pageParam
            : getPreviousPageParam(context.options, oldPages)
          promise = fetchPage(oldPages, manual, param, true)
        }

        // Refetch pages
        else {
          newPageParams = []

          const manual = typeof context.options.getNextPageParam === 'undefined'

          const shouldFetchFirstPage =
            refetchPage && oldPages[0]
              ? refetchPage(oldPages[0], 0, oldPages)
              : true

          // Fetch first page
          promise = shouldFetchFirstPage
            ? fetchPage([], manual, oldPageParams[0])
            : Promise.resolve(buildNewPages([], oldPageParams[0], oldPages[0]))

          // Fetch remaining pages
          for (let i = 1; i < oldPages.length; i++) {
            promise = promise.then(pages => {
              const shouldFetchNextPage =
                refetchPage && oldPages[i]
                  ? refetchPage(oldPages[i], i, oldPages)
                  : true

              if (shouldFetchNextPage) {
                const param = manual
                  ? oldPageParams[i]
                  : getNextPageParam(context.options, pages)
                return fetchPage(pages, manual, param)
              }
              return Promise.resolve(
                buildNewPages(pages, oldPageParams[i], oldPages[i])
              )
            })
          }
        }

        const finalPromise = promise.then(pages => ({
          pages,
          pageParams: newPageParams,
        }))

        context.signal?.addEventListener('abort', () => {
          cancelled = true
          abortController?.abort()
        })

        return finalPromise
      }
    },
  }
}

export function getNextPageParam(
  options: QueryOptions<any, any>,
  pages: unknown[]
): unknown | undefined {
  return options.getNextPageParam?.(pages[pages.length - 1], pages)
}

export function getPreviousPageParam(
  options: QueryOptions<any, any>,
  pages: unknown[]
): unknown | undefined {
  return options.getPreviousPageParam?.(pages[0], pages)
}

/**
 * Checks if there is a next page.
 * Returns `undefined` if it cannot be determined.
 */
export function hasNextPage(
  options: QueryOptions<any, any>,
  pages?: unknown
): boolean | undefined {
  if (options.getNextPageParam && Array.isArray(pages)) {
    const nextPageParam = getNextPageParam(options, pages)
    return (
      typeof nextPageParam !== 'undefined' &&
      nextPageParam !== null &&
      nextPageParam !== false
    )
  }
}

/**
 * Checks if there is a previous page.
 * Returns `undefined` if it cannot be determined.
 */
export function hasPreviousPage(
  options: QueryOptions<any, any>,
  pages?: unknown
): boolean | undefined {
  if (options.getPreviousPageParam && Array.isArray(pages)) {
    const previousPageParam = getPreviousPageParam(options, pages)
    return (
      typeof previousPageParam !== 'undefined' &&
      previousPageParam !== null &&
      previousPageParam !== false
    )
  }
}