omniai-web/src/stores/useGenerationStore.ts

import { create } from "zustand";
import type { WebGenerationPreviewTask } from "../types";

export type QueueItemStatus = "pending" | "running" | "completed" | "failed" | "cancelled";

export interface GenerationQueueItem {
  id: string;
  taskId?: string;
  title: string;
  type: "image" | "video" | "agent" | "digital-human" | "character-mix" | "ecommerce-video";
  status: QueueItemStatus;
  progress: number;
  prompt: string;
  createdAt: number;
  sourceView: string; // which page created this: "ecommerce", "workbench", "canvas", "agent"
  resultUrl?: string | null;
  error?: string | null;
  params?: Record<string, unknown>;
}

type PreviewTaskPatch = Partial<WebGenerationPreviewTask>;

interface PersistedQueueSnapshot {
  version: 1;
  items: GenerationQueueItem[];
  savedAt: number;
}

const STORAGE_KEY = "omniai:generation-queue";
const MAX_ITEMS = 80;
const STALE_MS = 2 * 60 * 60 * 1000; // 2 hours

function loadPersistedQueue(): GenerationQueueItem[] {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return [];
    const snapshot = JSON.parse(raw) as PersistedQueueSnapshot;
    if (Date.now() - (snapshot.savedAt || 0) > STALE_MS) {
      localStorage.removeItem(STORAGE_KEY);
      return [];
    }
    return snapshot.items.filter(
      (item) => item.status === "pending" || item.status === "running",
    );
  } catch {
    return [];
  }
}

function persistQueue(items: GenerationQueueItem[]): void {
  try {
    const snapshot: PersistedQueueSnapshot = { version: 1, items, savedAt: Date.now() };
    localStorage.setItem(STORAGE_KEY, JSON.stringify(snapshot));
  } catch { /* quota exceeded */ }
}

interface GenerationStoreState {
  queue: GenerationQueueItem[];
  tasks: WebGenerationPreviewTask[];
  addTask: (item: GenerationQueueItem) => void;
  updateTask: (id: string, patch: Partial<GenerationQueueItem>) => void;
  removeTask: (id: string) => void;
  setTasks: (tasks: WebGenerationPreviewTask[]) => void;
  appendTask: (task: WebGenerationPreviewTask) => void;
  mergeServerTasks: (serverTasks: WebGenerationPreviewTask[]) => void;
  clearTasks: () => void;
  getRunningTasks: () => GenerationQueueItem[];
  getPendingTasks: () => GenerationQueueItem[];
  getTasksByView: (sourceView: string) => GenerationQueueItem[];
  clearTerminal: () => void;
}

const initialQueue = loadPersistedQueue();

function trimTasks(tasks: WebGenerationPreviewTask[]): WebGenerationPreviewTask[] {
  return tasks.slice(0, MAX_ITEMS);
}

function mergePreviewTaskById(
  tasks: WebGenerationPreviewTask[],
  taskId: string | undefined,
  patch: PreviewTaskPatch,
): WebGenerationPreviewTask[] {
  if (!taskId) return tasks;
  let changed = false;
  const next = tasks.map((task) => {
    if (task.id !== taskId) return task;
    changed = true;
    return { ...task, ...patch };
  });
  return changed ? next : tasks;
}

function toPreviewTaskStatus(status: GenerationQueueItem["status"]): WebGenerationPreviewTask["status"] {
  if (status === "pending") return "queued";
  if (status === "cancelled") return "failed";
  return status;
}

function toPreviewTaskPatch(item: GenerationQueueItem): PreviewTaskPatch {
  const status = toPreviewTaskStatus(item.status);

  return {
    status,
    progress: item.status === "completed" ? 100 : item.progress,
    outputUrl: item.resultUrl || undefined,
    errorMessage: item.error || undefined,
  };
}

function toPreviewTask(item: GenerationQueueItem): WebGenerationPreviewTask | null {
  if (item.type === "ecommerce-video") return null;
  const type = item.type;
  const createdAt = Number.isFinite(item.createdAt)
    ? new Date(item.createdAt).toISOString()
    : new Date().toISOString();

  return {
    id: item.taskId || item.id,
    title: item.title,
    type,
    status: toPreviewTaskStatus(item.status),
    progress: item.status === "completed" ? 100 : item.progress,
    prompt: item.prompt,
    createdAt,
    projectId:
      typeof item.params?.projectId === "string" ? item.params.projectId : undefined,
    outputUrl: item.resultUrl || undefined,
    source: "preview",
    errorMessage: item.error || undefined,
  };
}

function upsertPreviewTask(
  tasks: WebGenerationPreviewTask[],
  task: WebGenerationPreviewTask | null,
): WebGenerationPreviewTask[] {
  if (!task) return tasks;
  return trimTasks([task, ...tasks.filter((item) => item.id !== task.id)]);
}

function previewTaskIdsForItem(item: GenerationQueueItem): string[] {
  return Array.from(new Set([item.taskId, item.id].filter(Boolean) as string[]));
}

export const useGenerationStore = create<GenerationStoreState>((set, get) => ({
  queue: initialQueue,
  tasks: [],

  addTask: (item) => {
    set((state) => {
      const next = [item, ...state.queue].slice(0, MAX_ITEMS);
      const previewTasks = upsertPreviewTask(state.tasks, toPreviewTask(item));
      persistQueue(next.filter((i) => i.status === "pending" || i.status === "running"));
      return { queue: next, tasks: previewTasks };
    });
  },

  updateTask: (id, patch) => {
    set((state) => {
      const next = state.queue.map((item) =>
        item.id === id ? { ...item, ...patch } : item,
      );
      const updated = next.find((item) => item.id === id);
      let previewTasks = state.tasks;
      if (updated) {
        const previewPatch = toPreviewTaskPatch(updated);
        for (const previewTaskId of previewTaskIdsForItem(updated)) {
          previewTasks = mergePreviewTaskById(previewTasks, previewTaskId, previewPatch);
        }
      }
      persistQueue(next.filter((i) => i.status === "pending" || i.status === "running"));
      return { queue: next, tasks: previewTasks };
    });
  },

  removeTask: (id) => {
    set((state) => {
      const next = state.queue.filter((item) => item.id !== id);
      persistQueue(next.filter((i) => i.status === "pending" || i.status === "running"));
      return { queue: next };
    });
  },

  setTasks: (tasks) => set({ tasks: trimTasks(tasks) }),

  appendTask: (task) => set((state) => ({
    tasks: trimTasks([task, ...state.tasks.filter((item) => item.id !== task.id)]),
  })),

  mergeServerTasks: (serverTasks) => set((state) => {
    const serverIds = new Set(serverTasks.map((task) => task.id));
    return {
      tasks: trimTasks([
        ...serverTasks,
        ...state.tasks.filter((task) => !serverIds.has(task.id)),
      ]),
    };
  }),

  clearTasks: () => {
    persistQueue([]);
    set({ tasks: [], queue: [] });
  },

  getRunningTasks: () => get().queue.filter((i) => i.status === "running" || i.status === "pending"),
  getPendingTasks: () => get().queue.filter((i) => i.status === "pending"),
  getTasksByView: (sourceView) => get().queue.filter((i) => i.sourceView === sourceView),

  clearTerminal: () => {
    set((state) => {
      const next = state.queue.filter(
        (i) => i.status === "pending" || i.status === "running",
      );
      persistQueue(next);
      return { queue: next };
    });
  },
}));