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 { saveGenerationRecord, type GenerationRecordStatus } from "../api/generationRecordClient"; import { startBackgroundPolling, subscribeToTaskUpdates, } from "../services/backgroundTaskRunner"; interface UseGenerationTasksOptions { sourceView: string; autoResume?: boolean; } type TerminalGenerationRecordStatus = Extract; function isTerminalStatus(status: GenerationQueueItem["status"] | undefined): status is TerminalGenerationRecordStatus { return status === "completed" || status === "failed" || status === "cancelled"; } function persistUnifiedTaskRecord(task: GenerationQueueItem): void { if (!isTerminalStatus(task.status)) return; void saveGenerationRecord({ clientRecordId: task.id, tool: task.sourceView, mode: task.type, title: task.title, status: task.status, prompt: task.prompt, taskIds: task.taskId ? [task.taskId] : [], assets: task.resultUrl ? [{ role: "result", mediaType: task.type === "video" || task.type === "ecommerce-video" ? "video" : "image", url: task.resultUrl, taskId: task.taskId, scope: `${task.sourceView}/result/${task.type}`, }] : [], config: task.params, result: { resultUrl: task.resultUrl, error: task.error, progress: task.progress, }, metadata: { queueCreatedAt: task.createdAt, source: "generation-queue", }, createdAt: new Date(task.createdAt).toISOString(), }); } 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) => { if (isTerminalStatus(patch.status)) { const current = queue.find((task) => task.id === id); if (current) persistUnifiedTaskRecord({ ...current, ...patch }); } updateStoredTask(id, patch); }, [queue, updateStoredTask], ); const markCompleted = useCallback( (id: string, resultUrl: string) => { const current = queue.find((task) => task.id === id); const patch: Partial = { status: "completed", progress: 100, resultUrl }; if (current) persistUnifiedTaskRecord({ ...current, ...patch }); updateStoredTask(id, patch); }, [queue, updateStoredTask], ); const markFailed = useCallback( (id: string, error: string) => { const current = queue.find((task) => task.id === id); const patch: Partial = { status: "failed", error }; if (current) persistUnifiedTaskRecord({ ...current, ...patch }); updateStoredTask(id, patch); }, [queue, 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, }; }