omniai-ds-code-package/src/hooks/useGenerationTasks.ts

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<GenerationRecordStatus, GenerationQueueItem["status"]>;

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<GenerationQueueItem, "id" | "createdAt">) => {
      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<GenerationQueueItem>) => {
      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<GenerationQueueItem> = { 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<GenerationQueueItem> = { 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,
  };
}