omniai-web/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 {
  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<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>) => {
      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,
  };
}