artcraft/frontend/apps/artcraft-webapp/src/pages/create-image/create-image-store.ts

import { create } from "zustand";
import { persist } from "zustand/middleware";
import type { RecreatePayload } from "../../lib/recreate";
import type { RefImage } from "../../components/prompt-box";

export interface GeneratedImage {
  media_token: string;
  cdn_url: string;
  maybe_thumbnail_template?: string;
}

export type ImageBatch = {
  id: string;
  prompt: string;
  status: "pending" | "complete" | "failed";
  images: GeneratedImage[];
  createdAt: number;
  requestedCount: number;
  modelLabel: string;
  jobToken?: string;
  failureReason?: string;
};

export type ImageUiState = {
  selectedModelId: string | null;
  prompt: string;
  aspectRatio: string;
  numImages: number;
  resolution: string | undefined;
  quality: string | undefined;
};

type CreateImageState = {
  batches: ImageBatch[];
  ui: ImageUiState;
  referenceImages: RefImage[];
  pendingRecreate: RecreatePayload | null;
  // Reference images sent from another page (library "Send to prompt").
  // Consumed by the create-image page, which applies the real per-model cap.
  pendingRefImages: RefImage[] | null;
  setUi: (patch: Partial<ImageUiState>) => void;
  setReferenceImages: (images: RefImage[]) => void;
  setPendingRecreate: (payload: RecreatePayload | null) => void;
  consumePendingRecreate: () => RecreatePayload | null;
  setPendingRefImages: (refs: RefImage[] | null) => void;
  consumePendingRefImages: () => RefImage[] | null;
  startBatch: (
    prompt: string,
    requestedCount: number,
    modelLabel: string,
  ) => string;
  setBatchJobToken: (batchId: string, jobToken: string) => void;
  completeBatch: (batchId: string, images: GeneratedImage[]) => void;
  failBatch: (batchId: string, reason?: string) => void;
  dismissBatch: (id: string) => void;
  clearCompleted: () => void;
  reset: () => void;
};

const DEFAULT_UI: ImageUiState = {
  selectedModelId: null,
  prompt: "",
  aspectRatio: "square",
  numImages: 1,
  resolution: undefined,
  quality: undefined,
};

export const useCreateImageStore = create<CreateImageState>()(
  persist(
    (set, get) => ({
      batches: [],
      ui: { ...DEFAULT_UI },
      referenceImages: [],
      pendingRecreate: null,
      pendingRefImages: null,

      setUi: (patch) =>
        set((s) => ({ ui: { ...s.ui, ...patch } })),

      setReferenceImages: (images) => set({ referenceImages: images }),

      setPendingRecreate: (payload) => set({ pendingRecreate: payload }),

      consumePendingRecreate: () => {
        const payload = get().pendingRecreate;
        if (payload) set({ pendingRecreate: null });
        return payload;
      },

      setPendingRefImages: (refs) => set({ pendingRefImages: refs }),

      consumePendingRefImages: () => {
        const refs = get().pendingRefImages;
        if (refs) set({ pendingRefImages: null });
        return refs;
      },

      startBatch: (prompt, requestedCount, modelLabel) => {
        const id = crypto.randomUUID();
        const batch: ImageBatch = {
          id,
          prompt,
          status: "pending",
          images: [],
          createdAt: Date.now(),
          requestedCount,
          modelLabel,
        };
        set((s) => ({ batches: [...s.batches, batch] }));
        return id;
      },

      setBatchJobToken: (batchId, jobToken) => {
        set((s) => ({
          batches: s.batches.map((b) =>
            b.id === batchId ? { ...b, jobToken } : b,
          ),
        }));
      },

      completeBatch: (batchId, images) => {
        set((s) => ({
          batches: s.batches.map((b) =>
            b.id === batchId
              ? { ...b, status: "complete" as const, images: images.slice(0, 4) }
              : b,
          ),
        }));
      },

      failBatch: (batchId, reason) => {
        set((s) => ({
          batches: s.batches.map((b) =>
            b.id === batchId
              ? { ...b, status: "failed" as const, failureReason: reason }
              : b,
          ),
        }));
      },

      dismissBatch: (id) => {
        set((s) => ({ batches: s.batches.filter((b) => b.id !== id) }));
      },

      clearCompleted: () => {
        set((s) => ({
          batches: s.batches.filter((b) => b.status !== "complete"),
        }));
      },

      reset: () => set({ batches: [] }),
    }),
    {
      name: "artcraft-image-prompt",
      // Persist prompt + lightweight generation settings so a full page reload
      // (e.g. returning from a credit top-up) keeps the user's draft.
      // Reference images are excluded: File handles aren't serializable and
      // uploaded images are base64 data URLs that can exceed the localStorage quota.
      partialize: (state) => ({ ui: state.ui }),
      merge: (persisted, current) => {
        const persistedUi = (persisted as { ui?: Partial<ImageUiState> } | null)?.ui;
        return { ...current, ui: { ...current.ui, ...(persistedUi ?? {}) } };
      },
    },
  ),
);