hyb-oa/client/src/hooks/useInvoiceOcrTracker.ts

import { useEffect, useMemo, useRef, useState } from "react";
import { useQuery } from "@tanstack/react-query";
import { subscribeAppEvents } from "../api/events";
import {
  fetchInvoiceStatuses,
  type InvoiceStatusItem,
} from "../api/invoicePool";

type UseInvoiceOcrTrackerOptions = {
  accessToken: string | null;
  enabled: boolean;
  pendingInvoiceIds: string[];
  setPendingInvoiceIds: (updater: string[] | ((previous: string[]) => string[])) => void;
  onAvailable: (invoiceIds: string[]) => void;
  onUnavailable: (items: InvoiceStatusItem[]) => void;
};

export function useInvoiceOcrTracker({
  accessToken,
  enabled,
  pendingInvoiceIds,
  setPendingInvoiceIds,
  onAvailable,
  onUnavailable,
}: UseInvoiceOcrTrackerOptions) {
  const [streamFailed, setStreamFailed] = useState(false);
  const firstPendingAtRef = useRef<number | null>(null);
  const handledIdsRef = useRef(new Set<string>());
  const pendingKey = pendingInvoiceIds.join(",");

  useEffect(() => {
    if (pendingInvoiceIds.length && firstPendingAtRef.current == null) {
      firstPendingAtRef.current = Date.now();
    }
    if (!pendingInvoiceIds.length) {
      firstPendingAtRef.current = null;
      handledIdsRef.current.clear();
    }
  }, [pendingInvoiceIds.length]);

  const statusQuery = useQuery({
    enabled: enabled && pendingInvoiceIds.length > 0,
    queryKey: ["invoice-pool", "status", pendingInvoiceIds],
    queryFn: () => fetchInvoiceStatuses(pendingInvoiceIds),
    refetchInterval: pendingInvoiceIds.length ? 3_000 : false,
  });

  useEffect(() => {
    if (!enabled || !accessToken || !pendingInvoiceIds.length) {
      return;
    }

    setStreamFailed(false);
    return subscribeAppEvents({
      accessToken,
      onEvent: (event) => {
        if (event.type !== "invoice.updated" || !event.invoiceId) {
          return;
        }
        if (pendingInvoiceIds.includes(event.invoiceId)) {
          void statusQuery.refetch();
        }
      },
      onError: () => setStreamFailed(true),
    });
  }, [accessToken, enabled, pendingKey]);

  useEffect(() => {
    const statuses = statusQuery.data?.items ?? [];
    if (!statuses.length) {
      return;
    }

    const availableIds = statuses
      .filter((item) => item.status === "available")
      .map((item) => item.id)
      .filter((invoiceId) => !handledIdsRef.current.has(invoiceId));
    const unavailableItems = statuses
      .filter((item) => ["ocr_failed", "archived", "bound"].includes(item.status))
      .filter((item) => !handledIdsRef.current.has(item.id));
    const finishedIds = new Set([...availableIds, ...unavailableItems.map((item) => item.id)]);

    if (!finishedIds.size) {
      return;
    }

    for (const invoiceId of finishedIds) {
      handledIdsRef.current.add(invoiceId);
    }

    if (availableIds.length) {
      onAvailable(availableIds);
    }
    if (unavailableItems.length) {
      onUnavailable(unavailableItems);
    }

    setPendingInvoiceIds((previous) => previous.filter((invoiceId) => !finishedIds.has(invoiceId)));
  }, [statusQuery.data?.items, onAvailable, onUnavailable, setPendingInvoiceIds]);

  const hasLongWait = Boolean(
    pendingInvoiceIds.length &&
      firstPendingAtRef.current &&
      Date.now() - firstPendingAtRef.current > 60_000,
  );

  const statusMessage = useMemo(() => {
    if (!pendingInvoiceIds.length) {
      return "";
    }

    if (hasLongWait) {
      return `${pendingInvoiceIds.length} 张发票仍在识别，系统会继续自动刷新，完成后会自动添加。`;
    }

    return `${pendingInvoiceIds.length} 张发票正在识别，正在自动刷新，完成后会自动添加。`;
  }, [hasLongWait, pendingInvoiceIds.length]);

  return {
    statusMessage,
    isTracking: pendingInvoiceIds.length > 0,
    streamFailed,
    refetch: statusQuery.refetch,
  };
}