import { useInfiniteQuery, useMutation, useQuery, useQueryClient, type InfiniteData, type QueryClient, } from "@tanstack/react-query" import { useEffect } from "react" import { getSelectionChart, getSelectionResultSectors, getSelectionResults, getSelectionRun, triggerSelectionRun, } from "./selection.api" import { selectionResultPageSize, type SelectionResults, type SelectionResultsQuery, type SelectionRunRequest, type SelectionSectors, type SelectionStrategy, } from "./selection.types" export type SelectionResultsListQuery = Omit 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", query.sector ?? "", ] as const export const selectionRunQueryKey = ( runId: string, query: SelectionResultsListQuery = defaultListQuery, ) => [ "selection", "runs", runId, query.pageSize, query.search ?? "", query.category ?? "all", query.sort ?? "code", query.sector ?? "", ] as const export const selectionSectorsQueryKey = ( strategy: SelectionStrategy, targetTradeDate?: string, ) => ["selection", "sectors", strategy, targetTradeDate] as const export function flattenSelectionResults( data: InfiniteData | 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, ReturnType, 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 useSelectionResultSectors( strategy: SelectionStrategy, targetTradeDate?: string, ) { return useQuery({ queryFn: ({ signal }) => getSelectionResultSectors(strategy, targetTradeDate, signal), queryKey: selectionSectorsQueryKey(strategy, targetTradeDate), }) } export function useSelectionRun( runId: string | null, resultQuery: SelectionResultsListQuery = defaultListQuery, ) { const queryClient = useQueryClient() const query = useInfiniteQuery< SelectionResults, Error, InfiniteData, ReturnType, 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 Promise.all([ queryClient.invalidateQueries({ queryKey: ["selection", "results", strategy], }), queryClient.invalidateQueries({ queryKey: ["selection", "sectors", 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 }