import { useMemo } from 'react'
import type { ModelProps } from '../types'
import { useSourceStore } from '../../providers/sources/store'
import { useShallow } from 'zustand/react/shallow'
import { buildQueryKey } from './build-query-key'

type QueryKeyArgs<TRequestOptions, TResponse> = Pick<
  ModelProps<TRequestOptions, TResponse>,
  never
> & { queryKeyPrefix: string[] }

type SingleSourceQueryArgs<TRequestOptions, TResponse> = QueryKeyArgs<
  TRequestOptions,
  TResponse
> & {
  sourceId: string | undefined
  sourcesId?: never
}

type MultiSourceQueryArgs<TRequestOptions, TResponse> = QueryKeyArgs<
  TRequestOptions,
  TResponse
> & {
  sourceId?: never
  sourcesId: string[]
}

// Overloads: when `sourceId` is provided return a single query key,
// when `sourcesId` is provided return an array of query keys
export function useQueryKey<TRequestOptions, TResponse>(
  args: SingleSourceQueryArgs<TRequestOptions, TResponse>,
): readonly unknown[]

export function useQueryKey<TRequestOptions, TResponse>(
  args: MultiSourceQueryArgs<TRequestOptions, TResponse>,
): unknown[][]

export function useQueryKey<TRequestOptions, TResponse>(
  args:
    | SingleSourceQueryArgs<TRequestOptions, TResponse>
    | MultiSourceQueryArgs<TRequestOptions, TResponse>,
): readonly unknown[] | unknown[][] {
  const { queryKeyPrefix, sourceId, sourcesId = [] } = args

  const sourcesArray = useMemo(() => {
    const array = sourceId !== undefined ? [sourceId] : sourcesId
    return array
  }, [sourceId, sourcesId])

  const sources = useSourceStore(
    useShallow((state) => sourcesArray.map((id) => state.sources[id])),
  )

  const queryKeys = useMemo(() => {
    return sourcesArray.map((currentSourceId, index) =>
      buildQueryKey(queryKeyPrefix, currentSourceId, sources[index]),
    )
  }, [queryKeyPrefix, sourcesArray, sources])

  if (sourceId !== undefined) {
    return queryKeys[0] as readonly unknown[]
  }

  return queryKeys
}
