zhixing-system/zhixing-web/src/features/selection/api/selection.query.ts

import {
  useInfiniteQuery,
  useMutation,
  useQuery,
  useQueryClient,
  type InfiniteData,
  type QueryClient,
} from "@tanstack/react-query"
import { useEffect } from "react"

import {
  getSelectionChart,
  getSelectionResults,
  getSelectionRun,
  triggerSelectionRun,
} from "./selection.api"
import {
  selectionResultPageSize,
  type SelectionResults,
  type SelectionResultsQuery,
  type SelectionRunRequest,
  type SelectionStrategy,
} from "./selection.types"

export type SelectionResultsListQuery = Omit<SelectionResultsQuery, "page">

const defaultListQuery: SelectionResultsListQuery = {
  pageSize: selectionResultPageSize,
}

export const selectionChartQueryKey = (
  tsCode?: string,
  targetTradeDate?: string,
  strategy?: SelectionStrategy,
) =>
  [
    "selection",
    "chart",
    tsCode ?? "none",
    targetTradeDate ?? "none",
    strategy ?? "none",
  ] as const

export const selectionResultsQueryKey = (
  strategy: SelectionStrategy,
  targetTradeDate?: string,
  query: SelectionResultsListQuery = defaultListQuery,
) =>
  [
    "selection",
    "results",
    strategy,
    targetTradeDate,
    query.pageSize,
    query.search ?? "",
    query.category ?? "all",
    query.sort ?? "code",
  ] as const

export const selectionRunQueryKey = (
  runId: string,
  query: SelectionResultsListQuery = defaultListQuery,
) =>
  [
    "selection",
    "runs",
    runId,
    query.pageSize,
    query.search ?? "",
    query.category ?? "all",
    query.sort ?? "code",
  ] as const

export function flattenSelectionResults(
  data: InfiniteData<SelectionResults> | undefined,
): SelectionResults | undefined {
  const firstPage = data?.pages[0]
  if (!firstPage) return undefined
  return {
    ...firstPage,
    stocks: data.pages.flatMap((page) => page.stocks),
  }
}

export function useSelectionResults(
  strategy: SelectionStrategy,
  targetTradeDate?: string,
  query: SelectionResultsListQuery = defaultListQuery,
) {
  return useInfiniteQuery<
    SelectionResults,
    Error,
    InfiniteData<SelectionResults>,
    ReturnType<typeof selectionResultsQueryKey>,
    number
  >({
    getNextPageParam: nextSelectionPageParam,
    initialPageParam: 1,
    queryFn: ({ pageParam, signal }) =>
      getSelectionResults(
        strategy,
        targetTradeDate,
        { ...query, page: pageParam },
        signal,
      ),
    queryKey: selectionResultsQueryKey(strategy, targetTradeDate, query),
  })
}

export function useSelectionChart(
  tsCode?: string,
  targetTradeDate?: string,
  strategy?: SelectionStrategy,
) {
  return useQuery({
    enabled: Boolean(tsCode && targetTradeDate),
    queryFn: ({ signal }) =>
      getSelectionChart(tsCode ?? "", targetTradeDate ?? "", strategy, signal),
    queryKey: selectionChartQueryKey(tsCode, targetTradeDate, strategy),
  })
}

export function useSelectionRun(
  runId: string | null,
  resultQuery: SelectionResultsListQuery = defaultListQuery,
) {
  const queryClient = useQueryClient()
  const query = useInfiniteQuery<
    SelectionResults,
    Error,
    InfiniteData<SelectionResults>,
    ReturnType<typeof selectionRunQueryKey>,
    number
  >({
    enabled: Boolean(runId),
    getNextPageParam: nextSelectionPageParam,
    initialPageParam: 1,
    queryFn: ({ pageParam, signal }) =>
      getSelectionRun(runId ?? "", { ...resultQuery, page: pageParam }, signal),
    queryKey: selectionRunQueryKey(runId ?? "none", resultQuery),
    refetchInterval: (currentQuery) =>
      currentQuery.state.data?.pages[0]?.status === "running" ? 1500 : false,
  })

  useEffect(() => {
    const result = query.data?.pages[0]
    if (!runId || !result || result.status === "running") return

    void invalidateSelectionResults(queryClient, result.strategy)
  }, [query.data, queryClient, runId])

  return query
}

export function useTriggerSelectionRun() {
  const queryClient = useQueryClient()

  return useMutation({
    mutationFn: ({
      request,
      signal,
    }: {
      request: SelectionRunRequest
      signal?: AbortSignal
    }) => triggerSelectionRun(request, signal),
    onSuccess: (_accepted, { request }) => {
      void invalidateSelectionResults(queryClient, request.strategy)
    },
  })
}

export function invalidateSelectionResults(
  queryClient: QueryClient,
  strategy: SelectionStrategy,
) {
  return queryClient.invalidateQueries({
    queryKey: ["selection", "results", strategy],
  })
}

function nextSelectionPageParam(
  lastPage: SelectionResults,
  pages: SelectionResults[],
  lastPageParam: number,
) {
  const loadedStockCount = pages.reduce(
    (count, page) => count + page.stocks.length,
    0,
  )
  const total = lastPage.stocks_total ?? lastPage.selected_stock_count
  if (
    (lastPage.status !== "success" && lastPage.status !== "partial_success") ||
    lastPage.stocks.length === 0 ||
    lastPage.page !== lastPageParam ||
    loadedStockCount >= total
  ) {
    return undefined
  }

  const nextPageParam = lastPageParam + 1
  return Number.isSafeInteger(nextPageParam) ? nextPageParam : undefined
}