import { useQuery, type UseQueryOptions } from '@tanstack/react-query'
import type { AddFilterOptions } from '@carto/api-client'
import type { ModelProps, ModelResponse } from '../types'
import type {
  GetModelQueryOptionsArgs,
  ModelQueryOptions,
} from './get-model-query-options'
import type { useSourceSlices } from './use-source-slices'

type SourceSlices = ReturnType<typeof useSourceSlices>

/**
 * Per-model query options factory signature. Each model file (categories.tsx,
 * formula.tsx, …) defines a function with this shape (`getCategories`,
 * `getFormula`, …) and feeds the same function back into `createModel` so the
 * hook's `useQuery` and any imperative `queryClient.prefetchQuery(getXxx(...))`
 * go through identical code — single source of truth for query-key + query-fn
 * construction.
 */
export type ModelQueryFactory<TWidgetProps extends object, TResponse> = (
  args: Omit<
    GetModelQueryOptionsArgs<TWidgetProps, TResponse>,
    'queryKeyPrefix' | 'getWidgetMethod'
  >,
) => ModelQueryOptions<TResponse>

interface CreateModelArgs<TWidgetProps extends object, TResponse> {
  buildQuery: ModelQueryFactory<TWidgetProps, TResponse>
}

export function createModel<TWidgetProps extends object, TResponse>({
  buildQuery,
}: CreateModelArgs<TWidgetProps, TResponse>) {
  function useModelHook<P extends ModelProps<TWidgetProps, TResponse>>(
    props: P,
    slices: SourceSlices,
  ): P extends { sourceId: string; sourcesId?: string[] }
    ? ModelResponse<TResponse>
    : ModelResponse<TResponse[]> {
    const {
      accessToken,
      sourceId,
      widgetProps,
      additionalFilter,
      useQueryProps: { enabled = true, ...useQueryProps } = {},
    } = props
    const { sourcesId, sources, filters } = slices

    const isSingleSource = !!sourceId

    // Imperative composition on every render. The `slices` argument is the
    // store snapshot read by the per-model wrapper via `useSourceSlices`; we
    // thread it into each `buildQuery` call as `sourceState` so the resulting
    // query key is a function of the *subscribed* values. That way a
    // `setSource(id, ...)` or filter mutation propagates: wrapper re-renders
    // → fresh `slices` → new `queryKey` → `useQuery` refetches.
    //
    // Skipping the previous `useMemo` here is intentional: the cost is one
    // queryKey array + one closure per render, and `useQuery` deep-compares
    // `queryKey` so it does not cause spurious refetches. An earlier version
    // tried a memo with `subscribedSources` / `subscribedFilters` "touched"
    // via `void` to keep them in deps, but `babel-plugin-react-compiler`
    // (used by this package's build) elided them — see
    // create-model-setsource-reactivity.test.tsx for the regression.
    const perSource = sourcesId.map((id, index) =>
      buildQuery({
        accessToken,
        sourceId: id,
        widgetProps,
        additionalFilter: pickAdditionalFilter(
          additionalFilter,
          index,
          isSingleSource,
        ),
        sourceIndex: index,
        totalSources: sourcesId.length,
        sourceState: {
          source: sources[index],
          filters: filters[index],
        },
      }),
    )

    const queryKey = perSource.flatMap((opt) => opt.queryKey)
    const allEnabled = perSource.every((opt) => opt.enabled)

    type QueryData = TResponse | TResponse[]
    const queryFn = async ({ signal }: { signal: AbortSignal }) => {
      const results = await Promise.allSettled(
        perSource.map((opt) => opt.queryFn({ signal })),
      )

      const data: TResponse[] = []
      let error: Error | undefined
      for (const [index, result] of results.entries()) {
        if (result.status === 'fulfilled') {
          data.push(result.value as TResponse)
          continue
        }
        error =
          result.reason instanceof Error
            ? result.reason
            : new Error(
                `Query failed for source ${sourcesId[index]}: ${String(result.reason)}`,
              )
        break
      }
      if (error) throw error
      return (isSingleSource ? data.at(0) : data) as QueryData
    }

    const response = useQuery<QueryData, Error>({
      queryKey,
      queryFn,
      enabled: allEnabled && enabled,
      ...useQueryProps,
    } as UseQueryOptions<QueryData, Error>)

    return response as P extends { sourceId: string; sourcesId?: string[] }
      ? ModelResponse<TResponse>
      : ModelResponse<TResponse[]>
  }

  return { useModelHook }
}

/**
 * For multi-source calls, `additionalFilter` may be `AddFilterOptions[]`
 * (apply to all) or `AddFilterOptions[][]` (per-source). For single source,
 * always `AddFilterOptions[]`. This selects the slice for the given source.
 */
function pickAdditionalFilter(
  additionalFilter: AddFilterOptions[] | AddFilterOptions[][] | undefined,
  index: number,
  isSingleSource: boolean,
): AddFilterOptions[] | undefined {
  if (!additionalFilter?.length) return undefined
  if (isSingleSource) return additionalFilter as AddFilterOptions[]
  if (Array.isArray(additionalFilter[0])) {
    return (additionalFilter as AddFilterOptions[][])[index]
  }
  return additionalFilter as AddFilterOptions[]
}
