omniai-web/src/features/workbench/workbenchStorage.ts

import {
  userKey,
  MESSAGE_STORAGE_KEY,
  ACTIVE_CONVERSATION_STORAGE_KEY,
  PROMPT_HISTORY_STORAGE_KEY,
  TASK_KEEPALIVE_STORAGE_KEY,
  WORKBENCH_TASK_STALE_MS,
  type ChatMessage,
  type WorkbenchKeepaliveTask,
} from "./workbenchConstants";
import { parseWorkbenchTimestampValue } from "./workbenchConstants";

export function readStoredMessages(): ChatMessage[] {
  if (typeof window === "undefined") return [];
  try {
    const raw = window.localStorage.getItem(userKey(MESSAGE_STORAGE_KEY));
    if (!raw) return [];
    const parsed = JSON.parse(raw) as unknown;
    if (!Array.isArray(parsed)) return [];
    return parsed.filter((item): item is ChatMessage => {
      if (!item || typeof item !== "object") return false;
      const candidate = item as Partial<ChatMessage>;
      return (
        typeof candidate.id === "string" &&
        (candidate.role === "user" || candidate.role === "assistant") &&
        typeof candidate.body === "string"
      );
    });
  } catch {
    return [];
  }
}

export function readStoredPromptHistory(): string[] {
  if (typeof window === "undefined") return [];
  try {
    const raw = window.localStorage.getItem(userKey(PROMPT_HISTORY_STORAGE_KEY));
    if (!raw) return [];
    const parsed = JSON.parse(raw) as unknown;
    return Array.isArray(parsed) ? parsed.filter((item): item is string => typeof item === "string") : [];
  } catch {
    return [];
  }
}

export function readStoredActiveConversationId(messages: ChatMessage[] = []): number | null {
  if (typeof window === "undefined") return null;
  try {
    const raw = window.localStorage.getItem(userKey(ACTIVE_CONVERSATION_STORAGE_KEY));
    const value = raw ? Number(raw) : NaN;
    if (Number.isFinite(value) && value > 0) return value;
  } catch {
    // Active conversation recovery is optional.
  }
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const candidate = messages[index]?.conversationId;
    if (typeof candidate === "number" && Number.isFinite(candidate) && candidate > 0) {
      return candidate;
    }
  }
  return null;
}

export function persistActiveConversationId(conversationId: number | null) {
  if (typeof window === "undefined") return;
  try {
    if (conversationId && Number.isFinite(conversationId)) {
      window.localStorage.setItem(userKey(ACTIVE_CONVERSATION_STORAGE_KEY), String(conversationId));
    } else {
      window.localStorage.removeItem(userKey(ACTIVE_CONVERSATION_STORAGE_KEY));
    }
  } catch {
    // Local history is a convenience; generation still works without it.
  }
}

export function persistMessages(messages: ChatMessage[]) {
  try {
    window.localStorage.setItem(userKey(MESSAGE_STORAGE_KEY), JSON.stringify(messages.slice(-60)));
  } catch {
    // Local history is a convenience; generation still works without it.
  }
}

export function clearWorkbenchLocalState() {
  if (typeof window === "undefined") return;
  try {
    window.localStorage.removeItem(userKey(MESSAGE_STORAGE_KEY));
    window.localStorage.removeItem(userKey(ACTIVE_CONVERSATION_STORAGE_KEY));
    window.localStorage.removeItem(userKey(TASK_KEEPALIVE_STORAGE_KEY));
  } catch {
    // Logout cleanup should never block the UI.
  }
}

export function persistPromptHistory(history: string[]) {
  try {
    window.localStorage.setItem(userKey(PROMPT_HISTORY_STORAGE_KEY), JSON.stringify(history.slice(0, 20)));
  } catch {
    // Local history is optional.
  }
}

// ─── Keepalive task persistence ──────────────────────────────────────

export function buildRecoverableTaskFromMessage(conversationId: number, message: ChatMessage): WorkbenchKeepaliveTask | null {
  if (message.role !== "assistant") return null;
  if (!(message.status === "thinking" && message.taskId && message.mode !== "chat")) return null;
  if (message.mode !== "image" && message.mode !== "video") return null;
  if (Date.now() - parseWorkbenchTimestampValue(message.createdAt) > WORKBENCH_TASK_STALE_MS) return null;

  const specs = message.result?.specs || [];
  return {
    taskId: message.taskId,
    conversationId,
    assistantMessageId: message.id,
    operation: message.taskStatusLabel?.includes("超分") ? "video-super-resolution" : "generation",
    mode: message.mode,
    modelLabel: specs[0] || message.author || message.mode,
    specs,
    referenceCount: message.attachments?.length || 0,
    progress: Math.max(10, Math.min(99, Number(message.taskProgress || 30))),
    statusLabel: message.taskStatusLabel || "任务恢复中...",
    startedAt: parseWorkbenchTimestampValue(message.createdAt) || Date.now(),
  };
}

export function readStoredKeepaliveTasks(): Record<string, WorkbenchKeepaliveTask> {
  if (typeof window === "undefined") return {};
  try {
    const raw = window.localStorage.getItem(userKey(TASK_KEEPALIVE_STORAGE_KEY));
    if (!raw) return {};
    const parsed = JSON.parse(raw) as unknown;
    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return {};

    const tasks: Record<string, WorkbenchKeepaliveTask> = {};
    Object.values(parsed as Record<string, Partial<WorkbenchKeepaliveTask>>).forEach((task) => {
      if (
        task &&
        typeof task.taskId === "string" &&
        typeof task.conversationId === "number" &&
        typeof task.assistantMessageId === "string" &&
        (task.mode === "image" || task.mode === "video")
      ) {
        tasks[task.taskId] = {
          taskId: task.taskId,
          conversationId: task.conversationId,
          assistantMessageId: task.assistantMessageId,
          concurrencySlotId: typeof task.concurrencySlotId === "string" ? task.concurrencySlotId : undefined,
          operation: task.operation === "video-super-resolution" ? "video-super-resolution" : "generation",
          mode: task.mode,
          modelLabel: task.modelLabel || task.mode,
          specs: Array.isArray(task.specs) ? task.specs.filter((item): item is string => typeof item === "string") : [],
          referenceCount: Number(task.referenceCount || 0),
          progress: Number(task.progress || 0),
          statusLabel: task.statusLabel || "Generating...",
          startedAt: Number(task.startedAt || Date.now()),
        };
      }
    });
    return tasks;
  } catch {
    return {};
  }
}

export function persistKeepaliveTasks(tasks: Record<string, WorkbenchKeepaliveTask>) {
  try {
    window.localStorage.setItem(userKey(TASK_KEEPALIVE_STORAGE_KEY), JSON.stringify(tasks));
  } catch {
    // Task restore is best-effort.
  }
}