omniai-web/src/services/backgroundTaskRunner.ts

import { useGenerationStore, type GenerationQueueItem } from "../stores/useGenerationStore";
import { aiGenerationClient } from "../api/aiGenerationClient";
import {
  buildLocalTimeoutMessage,
  buildTaskFailureInfo,
  getTaskTimeoutPolicy,
  isTaskLocallyTimedOut,
} from "../utils/taskLifecycle";

type PollCallback = (item: GenerationQueueItem) => void;

const activePollers = new Map<string, ReturnType<typeof setInterval>>();
const pollCallbacks = new Set<PollCallback>();

const POLL_INTERVAL = 3000;
const MAX_POLL_ATTEMPTS = 200; // Keep the previous 10-minute guard as a fallback.

export function subscribeToTaskUpdates(callback: PollCallback): () => void {
  pollCallbacks.add(callback);
  return () => { pollCallbacks.delete(callback); };
}

function notifyCallbacks(item: GenerationQueueItem): void {
  pollCallbacks.forEach((cb) => cb(item));
}

function getQueueItemKind(item: GenerationQueueItem): "image" | "video" | "text" {
  if (item.type === "image") return "image";
  if (item.type === "video" || item.type === "ecommerce-video") return "video";
  return "text";
}

function getQueueItemModel(item: GenerationQueueItem): string | undefined {
  return typeof item.params?.model === "string" ? item.params.model : undefined;
}

function pollTask(item: GenerationQueueItem, attemptsRef: { current: number }): void {
  const key = `poll-${item.id}`;
  if (activePollers.has(key)) return;

  const kind = getQueueItemKind(item);
  const timeoutPolicy = getTaskTimeoutPolicy({ kind, model: getQueueItemModel(item) });
  let lastProgress = Math.max(0, Number(item.progress || 0));
  let lastProgressAt = Date.now();

  const interval = setInterval(async () => {
    const current = useGenerationStore.getState().queue.find((i) => i.id === item.id);
    if (!current || current.status === "completed" || current.status === "failed" || current.status === "cancelled") {
      cleanupPoll(key);
      return;
    }

    attemptsRef.current++;
    const timeoutReason = isTaskLocallyTimedOut({
      startedAt: current.createdAt || item.createdAt || Date.now(),
      lastProgressAt,
      progress: lastProgress,
      policy: timeoutPolicy,
    });
    if (timeoutReason || attemptsRef.current > MAX_POLL_ATTEMPTS) {
      const error = buildLocalTimeoutMessage(kind);
      useGenerationStore.getState().updateTask(item.id, {
        status: "failed",
        error,
      });
      notifyCallbacks({ ...item, status: "failed", error });
      cleanupPoll(key);
      return;
    }

    try {
      const status = await aiGenerationClient.getTaskStatus(current.taskId || item.taskId || "");
      const nextProgress = Number(status.progress || 0);
      if (nextProgress > lastProgress || status.status === "completed") {
        lastProgress = Math.max(lastProgress, nextProgress);
        lastProgressAt = Date.now();
      }

      const patch: Partial<GenerationQueueItem> = {
        progress: status.progress,
        resultUrl: status.resultUrl || current.resultUrl,
        error: status.error || current.error,
      };

      if (status.status === "completed") {
        patch.status = "completed";
        useGenerationStore.getState().updateTask(item.id, patch);
        notifyCallbacks({ ...item, ...patch, status: "completed" });
        cleanupPoll(key);
      } else if (status.status === "failed" || status.status === "cancelled") {
        patch.status = "failed";
        patch.error = buildTaskFailureInfo(status.error).message;
        useGenerationStore.getState().updateTask(item.id, patch);
        notifyCallbacks({ ...item, ...patch, status: "failed" });
        cleanupPoll(key);
      } else {
        patch.status = "running";
        useGenerationStore.getState().updateTask(item.id, patch);
        notifyCallbacks({ ...item, ...patch, status: "running" });
      }
    } catch {
      // Network errors during polling are retried until the lifecycle guard trips.
    }
  }, POLL_INTERVAL);

  activePollers.set(key, interval);
}

function cleanupPoll(key: string): void {
  const interval = activePollers.get(key);
  if (interval) {
    clearInterval(interval);
    activePollers.delete(key);
  }
}

export function startBackgroundPolling(): void {
  const tasks = useGenerationStore.getState().getRunningTasks();
  const attemptsMap = new Map<string, { current: number }>();

  tasks.forEach((task) => {
    if (task.taskId) {
      if (!attemptsMap.has(task.id)) {
        attemptsMap.set(task.id, { current: 0 });
      }
      pollTask(task, attemptsMap.get(task.id)!);
    }
  });
}

export function resumeTaskPolling(taskId: string, storeId: string): void {
  const task = useGenerationStore.getState().queue.find((i) => i.id === storeId);
  if (task && task.status !== "completed" && task.status !== "failed") {
    pollTask(task, { current: 0 });
  }
}

export function stopAllPolling(): void {
  activePollers.forEach((interval) => clearInterval(interval));
  activePollers.clear();
}

export function recoverAndResumeTasks(): void {
  const pendingTasks = useGenerationStore.getState().getRunningTasks();
  if (!pendingTasks.length) return;

  pendingTasks.forEach((task) => {
    if (task.taskId) {
      useGenerationStore.getState().updateTask(task.id, { status: "pending" });
    } else {
      useGenerationStore.getState().updateTask(task.id, {
        status: "failed",
        error: "页面刷新后任务没有服务端 ID，已释放本地占用，请重新提交。",
      });
    }
  });

  setTimeout(() => startBackgroundPolling(), 500);
}