lovart3/hooks/use-pending-generation.ts

"use client";

import { useEffect, useState, useCallback } from "react";

/**
 * 持久化的生成任务状态
 */
export interface PendingGeneration {
  chatId: string;
  modelId: string;
  modelType: "gemini" | "volcano-vision" | "volcano-video" | "qianwen-video" | "alibaba-image" | "other";
  timestamp: number;
  userMessageId?: string;
}

const STORAGE_KEY = "ai-chatbot-pending-generations";
const EXPIRY_TIME = 5 * 60 * 1000; // 5 分钟超时

/**
 * 获取所有未完成的生成任务
 */
function getPendingGenerations(): PendingGeneration[] {
  if (typeof window === "undefined") return [];
  try {
    const data = localStorage.getItem(STORAGE_KEY);
    if (!data) return [];
    const generations: PendingGeneration[] = JSON.parse(data);
    // 过滤掉过期的任务
    const now = Date.now();
    return generations.filter((g) => now - g.timestamp < EXPIRY_TIME);
  } catch {
    return [];
  }
}

/**
 * 保存未完成的生成任务
 */
function savePendingGenerations(generations: PendingGeneration[]) {
  if (typeof window === "undefined") return;
  try {
    localStorage.setItem(STORAGE_KEY, JSON.stringify(generations));
  } catch (e) {
    console.error("Failed to save pending generations:", e);
  }
}

/**
 * Hook: 管理持久化的生成任务状态
 * 解决用户离开页面再返回时加载状态丢失的问题
 */
export function usePendingGeneration(chatId: string) {
  const [pendingGeneration, setPendingGeneration] = useState<PendingGeneration | null>(null);

  // 初始化：检查是否有当前 chatId 的未完成任务
  useEffect(() => {
    const generations = getPendingGenerations();
    console.log("[PendingGeneration] 初始化检查, chatId:", chatId);
    console.log("[PendingGeneration] localStorage 中的所有任务:", generations);
    
    const found = generations.find((g) => g.chatId === chatId);
    if (found) {
      setPendingGeneration(found);
      console.log("[PendingGeneration] ✅ 恢复处理中状态:", found);
    } else {
      console.log("[PendingGeneration] ❌ 未找到匹配的任务");
    }
  }, [chatId]);

  // 开始生成：保存状态
  const startGeneration = useCallback(
    (modelId: string, modelType: PendingGeneration["modelType"], userMessageId?: string) => {
      const newGeneration: PendingGeneration = {
        chatId,
        modelId,
        modelType,
        timestamp: Date.now(),
        userMessageId,
      };

      // 更新 localStorage
      const generations = getPendingGenerations();
      // 移除同一 chatId 的旧记录
      const filtered = generations.filter((g) => g.chatId !== chatId);
      savePendingGenerations([...filtered, newGeneration]);

      setPendingGeneration(newGeneration);
      console.log("[PendingGeneration] 开始生成:", newGeneration);
    },
    [chatId]
  );

  // 完成生成：清除状态
  const completeGeneration = useCallback(() => {
    const generations = getPendingGenerations();
    const filtered = generations.filter((g) => g.chatId !== chatId);
    savePendingGenerations(filtered);

    setPendingGeneration(null);
    console.log("[PendingGeneration] 生成完成，清除状态");
  }, [chatId]);

  // 检查是否有恢复的任务（用户离开再返回）
  const hasRestoredPending = pendingGeneration !== null && Date.now() - pendingGeneration.timestamp > 1000;

  return {
    pendingGeneration,
    hasRestoredPending,
    startGeneration,
    completeGeneration,
  };
}