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

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

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

export const selectionResultsQueryKey = (
  strategy: SelectionStrategy,
  targetTradeDate?: string,
  query: SelectionResultsQuery = { page: 1, pageSize: 10 },
) =>
  [
    "selection",
    "results",
    strategy,
    targetTradeDate,
    query.page,
    query.pageSize,
    query.search ?? "",
    query.category ?? "all",
  ] as const

export const selectionRunQueryKey = (
  runId: string,
  query: SelectionResultsQuery = { page: 1, pageSize: 10 },
) =>
  [
    "selection",
    "runs",
    runId,
    query.page,
    query.pageSize,
    query.search ?? "",
    query.category ?? "all",
  ] as const

export function useSelectionResults(
  strategy: SelectionStrategy,
  targetTradeDate?: string,
  query: SelectionResultsQuery = { page: 1, pageSize: 10 },
) {
  return useQuery({
    queryFn: ({ signal }) =>
      getSelectionResults(strategy, targetTradeDate, query, signal),
    queryKey: selectionResultsQueryKey(strategy, targetTradeDate, query),
  })
}

export function useSelectionRun(
  runId: string | null,
  resultQuery: SelectionResultsQuery = { page: 1, pageSize: 10 },
) {
  const queryClient = useQueryClient()
  const query = useQuery({
    enabled: Boolean(runId),
    queryFn: ({ signal }) => getSelectionRun(runId ?? "", resultQuery, signal),
    queryKey: selectionRunQueryKey(runId ?? "none", resultQuery),
    refetchInterval: (query) =>
      query.state.data?.status === "running" ? 1500 : false,
  })

  useEffect(() => {
    const result = query.data
    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],
  })
}