MiragenFlow/web/src/stores/use-asset-store.ts

"use client";

import { create } from "zustand";
import { persist, type PersistStorage, type StorageValue } from "zustand/middleware";

import { nanoid } from "nanoid";
import { localForageStorage } from "@/lib/localforage-storage";
import { cleanupUnusedImages, resolveImageUrl, uploadImage } from "@/services/image-storage";

export type AssetKind = "text" | "image";
export type TextAsset = AssetBase<"text"> & { data: { content: string } };
export type ImageAsset = AssetBase<"image"> & { data: { dataUrl: string; storageKey?: string; width: number; height: number; bytes: number; mimeType: string } };
export type Asset = TextAsset | ImageAsset;

type AssetBase<T extends AssetKind> = {
  id: string;
  kind: T;
  title: string;
  coverUrl: string;
  tags: string[];
  source?: string;
  note?: string;
  createdAt: string;
  updatedAt: string;
  metadata?: Record<string, unknown>;
};

type AssetStore = {
  assets: Asset[];
  addAsset: (asset: Omit<Asset, "id" | "createdAt" | "updatedAt">) => string;
  updateAsset: (id: string, patch: Partial<Omit<Asset, "id" | "createdAt">>) => void;
  removeAsset: (id: string) => void;
  cleanupImages: (extra?: unknown) => void;
};

const ASSET_STORE_KEY = "infinite-canvas:asset_store";

const assetStorage: PersistStorage<AssetStore> = {
  getItem: async (name) => {
    const value = await localForageStorage.getItem(name);
    if (!value) return null;
    const parsed = JSON.parse(value) as StorageValue<AssetStore>;
    parsed.state.assets = await Promise.all(parsed.state.assets.map(async (asset) => {
      if (asset.kind !== "image") return asset;
      if (asset.data.storageKey) return {
        ...asset,
        coverUrl: asset.coverUrl.startsWith("blob:") ? await resolveImageUrl(asset.data.storageKey, asset.coverUrl) : asset.coverUrl,
        data: { ...asset.data, dataUrl: await resolveImageUrl(asset.data.storageKey, asset.data.dataUrl) },
      };
      if (!asset.data.dataUrl.startsWith("data:image/")) return asset;
      const image = await uploadImage(asset.data.dataUrl);
      return { ...asset, coverUrl: asset.coverUrl.startsWith("data:image/") ? image.url : asset.coverUrl, data: { ...asset.data, dataUrl: image.url, storageKey: image.storageKey, bytes: image.bytes, mimeType: image.mimeType } };
    }));
    return parsed;
  },
  setItem: (name, value) => localForageStorage.setItem(name, JSON.stringify(value)),
  removeItem: (name) => localForageStorage.removeItem(name),
};

export const useAssetStore = create<AssetStore>()(
  persist(
    (set, get) => ({
      assets: [],
      addAsset: (asset) => {
        const now = new Date().toISOString();
        const id = nanoid();
        set((state) => ({ assets: [{ ...asset, id, createdAt: now, updatedAt: now } as Asset, ...state.assets] }));
        return id;
      },
      updateAsset: (id, patch) => set((state) => ({
        assets: state.assets.map((asset) => asset.id === id ? { ...asset, ...patch, updatedAt: new Date().toISOString() } as Asset : asset),
      })),
      removeAsset: (id) => set((state) => {
        const assets = state.assets.filter((asset) => asset.id !== id);
        get().cleanupImages({ assets });
        return { assets };
      }),
      cleanupImages: (extra) => {
        window.setTimeout(async () => {
          const { useCanvasStore } = await import("@/app/(user)/canvas/stores/use-canvas-store");
          await cleanupUnusedImages({ assets: get().assets, projects: useCanvasStore.getState().projects, extra });
        }, 0);
      },
    }),
    {
      name: ASSET_STORE_KEY,
      storage: assetStorage,
      partialize: (state) => ({ assets: state.assets }) as StorageValue<AssetStore>["state"],
    },
  ),
);