omniai-web/src/features/canvas/useCanvasGeneration.ts

import { type Dispatch, type SetStateAction, useCallback, useEffect, useRef, useState } from "react";
import type {
  CanvasImageGenerationState,
  CanvasImageNode,
  CanvasTextGenerationState,
  CanvasVideoGenerationState,
  CanvasVideoNode,
} from "./canvasTypes";
import { aiGenerationClient } from "../../api/aiGenerationClient";
import { createCanvasAssetRefFromGeneratedResult, persistCanvasGeneratedResultAsset } from "./canvasAssetPersistence";
import { waitForImageTaskResult, waitForVideoTaskResult } from "./canvasUtils";

const CANVAS_GEN_KEEPALIVE_KEY = "omniai:canvas-gen-tasks";

interface CanvasGenKeepaliveEntry {
  taskId: string;
  nodeId: string;
  nodeKind: "image" | "video";
  projectId: string;
}

function loadCanvasGenKeepalive(): CanvasGenKeepaliveEntry[] {
  try {
    const raw = window.localStorage.getItem(CANVAS_GEN_KEEPALIVE_KEY);
    return raw ? JSON.parse(raw) : [];
  } catch {
    return [];
  }
}

function saveCanvasGenKeepalive(entries: CanvasGenKeepaliveEntry[]): void {
  window.localStorage.setItem(CANVAS_GEN_KEEPALIVE_KEY, JSON.stringify(entries));
}

export function addCanvasGenKeepalive(taskId: string, nodeId: string, nodeKind: "image" | "video", projectId: string): void {
  const entries = loadCanvasGenKeepalive().filter((e) => e.taskId !== taskId);
  entries.push({ taskId, nodeId, nodeKind, projectId });
  saveCanvasGenKeepalive(entries);
}

export function removeCanvasGenKeepalive(taskId: string): void {
  saveCanvasGenKeepalive(loadCanvasGenKeepalive().filter((e) => e.taskId !== taskId));
}

export function cancelCanvasGenKeepaliveOnUnload(): void {
  const entries = loadCanvasGenKeepalive();
  if (!entries.length) return;
  entries.forEach((entry) => aiGenerationClient.cancelTaskOnUnload(entry.taskId));
  saveCanvasGenKeepalive([]);
}

export interface UseCanvasGenerationParams {
  setImageNodes: Dispatch<SetStateAction<CanvasImageNode[]>>;
  setVideoNodes: Dispatch<SetStateAction<CanvasVideoNode[]>>;
}

export function useCanvasGeneration(params: UseCanvasGenerationParams) {
  const { setImageNodes, setVideoNodes } = params;

  const [textGenerationState, setTextGenerationState] = useState<Record<string, CanvasTextGenerationState>>({});
  const [imageGenerationState, setImageGenerationState] = useState<Record<string, CanvasImageGenerationState>>({});
  const [videoGenerationState, setVideoGenerationState] = useState<Record<string, CanvasVideoGenerationState>>({});
  const [generationToast, setGenerationToast] = useState<string | null>(null);

  const imageGenerationInFlightRef = useRef(new Set<string>());
  const videoGenerationInFlightRef = useRef(new Set<string>());
  const textGenerationInFlightRef = useRef(new Set<string>());
  const textGenerationAbortControllersRef = useRef(new Map<string, AbortController>());
  const imageGenerationAbortRef = useRef(new Map<string, { current: boolean }>());
  const videoGenerationAbortRef = useRef(new Map<string, { current: boolean }>());
  const canvasGenKeepaliveRestoredRef = useRef(false);

  const setTextGenerationStatus = (nodeId: string, state: CanvasTextGenerationState) => {
    setTextGenerationState((current) => ({ ...current, [nodeId]: state }));
  };

  const setImageGenerationStatus = (nodeId: string, state: CanvasImageGenerationState) => {
    setImageGenerationState((current) => ({ ...current, [nodeId]: state }));
  };

  const setVideoGenerationStatus = (nodeId: string, state: CanvasVideoGenerationState) => {
    setVideoGenerationState((current) => ({ ...current, [nodeId]: state }));
  };

  const abortAllGenerationPollers = useCallback(() => {
    textGenerationAbortControllersRef.current.forEach((c) => c.abort());
    textGenerationAbortControllersRef.current.clear();
    imageGenerationAbortRef.current.forEach((ref) => { ref.current = true; });
    imageGenerationAbortRef.current.clear();
    videoGenerationAbortRef.current.forEach((ref) => { ref.current = true; });
    videoGenerationAbortRef.current.clear();
  }, []);

  // Toast auto-dismiss
  useEffect(() => {
    if (!generationToast) return undefined;
    const timer = window.setTimeout(() => setGenerationToast(null), 1800);
    return () => window.clearTimeout(timer);
  }, [generationToast]);

  const restoreKeepaliveTasks = (
    projectId: string,
    imageNodesList: CanvasImageNode[],
    videoNodesList: CanvasVideoNode[],
  ) => {
    if (canvasGenKeepaliveRestoredRef.current) return;
    canvasGenKeepaliveRestoredRef.current = true;
    const storedTasks = loadCanvasGenKeepalive().filter((e) => e.projectId === projectId);
    for (const entry of storedTasks) {
      const nodeExists = [...imageNodesList, ...videoNodesList].some((n) => n.id === entry.nodeId);
      if (!nodeExists) {
        removeCanvasGenKeepalive(entry.taskId);
        continue;
      }
      if (entry.nodeKind === "image") {
        imageGenerationInFlightRef.current.add(entry.nodeId);
        const abortRef = { current: false };
        imageGenerationAbortRef.current.set(entry.nodeId, abortRef);
        setImageGenerationStatus(entry.nodeId, { status: "running", message: "正在恢复图片生成", progress: 20 });
        void waitForImageTaskResult(entry.taskId, (status) => {
          const progress = Math.max(18, Math.min(status.status === "completed" ? 100 : 96, Math.trunc(status.progress || 0)));
          setImageGenerationStatus(entry.nodeId, { status: "running", message: "图片生成中", progress });
        }, abortRef).then(async (outputUrl) => {
          if (abortRef.current || !outputUrl) return;
          removeCanvasGenKeepalive(entry.taskId);
          setImageGenerationStatus(entry.nodeId, { status: "success", message: "生成完成", progress: 100 });
          const ref = createCanvasAssetRefFromGeneratedResult({
            url: outputUrl, mediaType: "image/png", resultType: "image", taskId: entry.taskId, originalUrl: outputUrl,
          });
          setImageNodes((current) => current.map((n) =>
            n.id === entry.nodeId ? { ...n, imageUrl: outputUrl, assetRef: ref, taskRef: { taskId: entry.taskId, status: "completed", resultUrl: outputUrl, updatedAt: new Date().toISOString() } } : n
          ));
          let durableRef = ref;
          try {
            durableRef = await persistCanvasGeneratedResultAsset({
              title: "image-node", url: outputUrl, mediaType: "image/png", resultType: "image", taskId: entry.taskId, originalUrl: outputUrl,
            });
          } catch { /* keep temp URL on failure */ }
          if (durableRef.url !== outputUrl || durableRef.ossKey) {
            setImageNodes((current) => current.map((n) =>
              n.id === entry.nodeId ? { ...n, imageUrl: durableRef.url, assetRef: durableRef, taskRef: { taskId: entry.taskId, status: "completed", resultUrl: durableRef.url, updatedAt: new Date().toISOString() } } : n
            ));
          }
        }).catch(() => {
          if (abortRef.current) return;
          removeCanvasGenKeepalive(entry.taskId);
          setImageGenerationStatus(entry.nodeId, { status: "error", message: "图片生成失败" });
        }).finally(() => {
          imageGenerationInFlightRef.current.delete(entry.nodeId);
          imageGenerationAbortRef.current.delete(entry.nodeId);
        });
      } else if (entry.nodeKind === "video") {
        videoGenerationInFlightRef.current.add(entry.nodeId);
        const abortRef = { current: false };
        videoGenerationAbortRef.current.set(entry.nodeId, abortRef);
        setVideoGenerationStatus(entry.nodeId, { status: "running", message: "正在恢复视频生成", progress: 20 });
        void waitForVideoTaskResult(entry.taskId, (status) => {
          const progress = Math.max(18, Math.min(status.status === "completed" ? 100 : 96, Math.trunc(status.progress || 0)));
          setVideoGenerationStatus(entry.nodeId, { status: "running", message: "视频生成中", progress });
        }, abortRef).then(async (outputUrl) => {
          if (abortRef.current || !outputUrl) return;
          removeCanvasGenKeepalive(entry.taskId);
          setVideoGenerationStatus(entry.nodeId, { status: "success", message: "生成完成", progress: 100 });
          const ref = createCanvasAssetRefFromGeneratedResult({
            url: outputUrl, mediaType: "video/mp4", resultType: "video", taskId: entry.taskId, originalUrl: outputUrl,
          });
          setVideoNodes((current) => current.map((n) =>
            n.id === entry.nodeId ? { ...n, videoUrl: outputUrl, assetRef: ref, taskRef: { taskId: entry.taskId, status: "completed", resultUrl: outputUrl, updatedAt: new Date().toISOString() } } : n
          ));
          let durableRef = ref;
          try {
            durableRef = await persistCanvasGeneratedResultAsset({
              title: "video-node", url: outputUrl, mediaType: "video/mp4", resultType: "video", taskId: entry.taskId, originalUrl: outputUrl,
            });
          } catch { /* keep temp URL on failure */ }
          if (durableRef.url !== outputUrl || durableRef.ossKey) {
            setVideoNodes((current) => current.map((n) =>
              n.id === entry.nodeId ? { ...n, videoUrl: durableRef.url, assetRef: durableRef, taskRef: { taskId: entry.taskId, status: "completed", resultUrl: durableRef.url, updatedAt: new Date().toISOString() } } : n
            ));
          }
        }).catch(() => {
          if (abortRef.current) return;
          removeCanvasGenKeepalive(entry.taskId);
          setVideoGenerationStatus(entry.nodeId, { status: "error", message: "视频生成失败" });
        }).finally(() => {
          videoGenerationInFlightRef.current.delete(entry.nodeId);
          videoGenerationAbortRef.current.delete(entry.nodeId);
        });
      }
    }
  };

  const resetGenerationState = () => {
    abortAllGenerationPollers();
    textGenerationInFlightRef.current.clear();
    imageGenerationInFlightRef.current.clear();
    videoGenerationInFlightRef.current.clear();
    setTextGenerationState({});
    setImageGenerationState({});
    setVideoGenerationState({});
  };

  // Stop all in-flight front-end polling/setState when the canvas unmounts (route change).
  // Keepalive records are intentionally preserved so restoreKeepaliveTasks can resume on return.
  useEffect(() => {
    return () => {
      abortAllGenerationPollers();
    };
  }, [abortAllGenerationPollers]);

  useEffect(() => {
    const handlePageHide = () => {
      cancelCanvasGenKeepaliveOnUnload();
      abortAllGenerationPollers();
      textGenerationInFlightRef.current.clear();
      imageGenerationInFlightRef.current.clear();
      videoGenerationInFlightRef.current.clear();
      setTextGenerationState({});
      setImageGenerationState({});
      setVideoGenerationState({});
    };
    const handleOnline = () => {
      aiGenerationClient.flushPendingTaskCancellations();
    };
    window.addEventListener("pagehide", handlePageHide);
    window.addEventListener("online", handleOnline);
    aiGenerationClient.flushPendingTaskCancellations();
    return () => {
      window.removeEventListener("pagehide", handlePageHide);
      window.removeEventListener("online", handleOnline);
    };
  }, [abortAllGenerationPollers]);

  return {
    textGenerationState,
    imageGenerationState,
    videoGenerationState,
    generationToast,
    setGenerationToast,
    imageGenerationInFlightRef,
    videoGenerationInFlightRef,
    textGenerationInFlightRef,
    textGenerationAbortControllersRef,
    imageGenerationAbortRef,
    videoGenerationAbortRef,
    canvasGenKeepaliveRestoredRef,
    setTextGenerationStatus,
    setImageGenerationStatus,
    setVideoGenerationStatus,
    restoreKeepaliveTasks,
    resetGenerationState,
  };
}