import { useEffect, useMemo, useRef, useCallback } from "react"; import { useShallow } from "zustand/react/shallow"; import type { GenerationQueueItem } from "../stores/useGenerationStore"; import { useGenerationStore } from "../stores/useGenerationStore"; import { startBackgroundPolling, subscribeToTaskUpdates, } from "../services/backgroundTaskRunner"; interface UseGenerationTasksOptions { sourceView: string; autoResume?: boolean; } export function useGenerationTasks(options: UseGenerationTasksOptions) { const { sourceView, autoResume = true } = options; const { queue, addTask, updateTask: updateStoredTask, getRunningTasks, } = useGenerationStore(useShallow((s) => ({ queue: s.queue, addTask: s.addTask, updateTask: s.updateTask, getRunningTasks: s.getRunningTasks, }))); const pollingStartedRef = useRef(false); // ── Auto-resume: re-subscribe to running tasks on mount ──── useEffect(() => { if (!autoResume || pollingStartedRef.current) return; pollingStartedRef.current = true; const active = getRunningTasks().filter((t) => t.sourceView === sourceView); if (active.length > 0) { startBackgroundPolling(); } return () => { pollingStartedRef.current = false; }; }, [autoResume, sourceView, getRunningTasks]); // ── Subscribe to live updates ─────────────────────────── useEffect(() => { return subscribeToTaskUpdates((updated) => { updateStoredTask(updated.id, updated); }); }, [updateStoredTask]); // ── View-scoped computed lists ────────────────────────── const myTasks = useMemo( () => queue.filter((t) => t.sourceView === sourceView), [queue, sourceView], ); const activeTasks = useMemo( () => myTasks.filter((t) => t.status === "running" || t.status === "pending"), [myTasks], ); const completedTasks = useMemo( () => myTasks.filter((t) => t.status === "completed"), [myTasks], ); const failedTasks = useMemo( () => myTasks.filter((t) => t.status === "failed"), [myTasks], ); // ── Actions ───────────────────────────────────────────── const submitTask = useCallback( (task: Omit) => { const id = `gen-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; addTask({ ...task, id, createdAt: Date.now() }); return id; }, [addTask], ); const updateTask = useCallback( (id: string, patch: Partial) => { updateStoredTask(id, patch); }, [updateStoredTask], ); const markCompleted = useCallback( (id: string, resultUrl: string) => { updateStoredTask(id, { status: "completed", progress: 100, resultUrl }); }, [updateStoredTask], ); const markFailed = useCallback( (id: string, error: string) => { updateStoredTask(id, { status: "failed", error }); }, [updateStoredTask], ); const retryTask = useCallback( (id: string) => { const task = queue.find((t) => t.id === id); if (task) { updateStoredTask(id, { status: "pending", progress: 0, error: null }); } }, [queue, updateStoredTask], ); return { tasks: myTasks, activeTasks, completedTasks, failedTasks, submitTask, updateTask, markCompleted, markFailed, retryTask, hasActiveTasks: activeTasks.length > 0, }; }