stable-diffusion.cpp-rest/webui/lib/storage.ts

import { useState } from 'react';

// Safe localStorage hook that handles quota errors
export function useLocalStorage<T>(
  key: string,
  initialValue: T,
  options?: {
    excludeLargeData?: boolean;
    maxSize?: number; // in bytes
  }
) {
  const { excludeLargeData = true, maxSize = 1024 * 1024 } = options || {}; // 1MB default

  // Get stored value
  const getStoredValue = (): T => {
    if (typeof window === 'undefined') {
      return initialValue;
    }

    try {
      const item = window.localStorage.getItem(key);
      if (item === null) {
        return initialValue;
      }
      return JSON.parse(item);
    } catch (error) {
      console.warn(`Error reading localStorage key "${key}":`, error);
      
      // If there's an error, clear the key and return initial value
      try {
        window.localStorage.removeItem(key);
      } catch (clearError) {
        console.warn(`Error clearing localStorage key "${key}":`, clearError);
      }
      
      return initialValue;
    }
  };

  const [storedValue, setStoredValue] = useState<T>(getStoredValue);

  // Set value to localStorage
  const setValue = (value: T | ((val: T) => T)) => {
    try {
      const valueToStore = value instanceof Function ? value(storedValue) : value;
      
      // Check if we should exclude large data
      if (excludeLargeData) {
        const serialized = JSON.stringify(valueToStore);
        
        // Check for large base64 images or data URLs
        if (serialized.includes('data:image/') || serialized.length > maxSize) {
          console.warn(`Skipping localStorage save for "${key}" - data too large or contains images`);
          setStoredValue(valueToStore);
          return;
        }
      }

      setStoredValue(valueToStore);
      
      if (typeof window !== 'undefined') {
        window.localStorage.setItem(key, JSON.stringify(valueToStore));
      }
    } catch (error) {
      console.warn(`Error saving localStorage key "${key}":`, error);
      
      // Still update the state even if localStorage fails
      const valueToStore = value instanceof Function ? value(storedValue) : value;
      setStoredValue(valueToStore);
      
      // Try to clear some space if it's a quota error
      if (error instanceof Error && error.name === 'QuotaExceededError') {
        try {
          // Clear some non-essential keys
          const keysToClear = ['img2img-form-data', 'text2img-form-data', 'inpainting-form-data'];
          keysToClear.forEach(k => {
            if (k !== key) {
              window.localStorage.removeItem(k);
            }
          });
        } catch (clearError) {
          console.warn('Error clearing localStorage space:', clearError);
        }
      }
    }
  };

  return [storedValue, setValue] as const;
}

// Session storage alternative for large data
export function useSessionStorage<T>(key: string, initialValue: T) {
  const [storedValue, setStoredValue] = useState<T>(() => {
    if (typeof window === 'undefined') {
      return initialValue;
    }

    try {
      const item = window.sessionStorage.getItem(key);
      return item ? JSON.parse(item) : initialValue;
    } catch (error) {
      console.warn(`Error reading sessionStorage key "${key}":`, error);
      return initialValue;
    }
  });

  const setValue = (value: T | ((val: T) => T)) => {
    try {
      const valueToStore = value instanceof Function ? value(storedValue) : value;
      setStoredValue(valueToStore);
      
      if (typeof window !== 'undefined') {
        window.sessionStorage.setItem(key, JSON.stringify(valueToStore));
      }
    } catch (error) {
      console.warn(`Error saving sessionStorage key "${key}":`, error);
      // Still update the state even if sessionStorage fails
      const valueToStore = value instanceof Function ? value(storedValue) : value;
      setStoredValue(valueToStore);
    }
  };

  return [storedValue, setValue] as const;
}

// Memory-only storage for large data
export function useMemoryStorage<T>(initialValue: T) {
  const [storedValue, setStoredValue] = useState<T>(initialValue);

  const setValue = (value: T | ((val: T) => T)) => {
    const valueToStore = value instanceof Function ? value(storedValue) : value;
    setStoredValue(valueToStore);
  };

  return [storedValue, setValue] as const;
}

// Storage for generated images with session persistence
interface StoredImage {
  url: string;
  timestamp: number;
  pageType: 'text2img' | 'img2img' | 'upscaler' | 'inpainting';
  jobId?: string;
}

interface GeneratedImagesState {
  images: StoredImage[];
  lastJobId?: string;
}

export function useGeneratedImages(pageType: 'text2img' | 'img2img' | 'upscaler' | 'inpainting') {
  const storageKey = `generated-images-${pageType}`;
  
  const [state, setState] = useSessionStorage<GeneratedImagesState>(
    storageKey,
    { images: [] }
  );

  // Clean up old images (older than 1 hour) on mount
  useState(() => {
    const now = Date.now();
    const oneHour = 60 * 60 * 1000;
    
    setState((prevState) => ({
      ...prevState,
      images: prevState.images.filter(img => now - img.timestamp < oneHour)
    }));
  });

  const addImages = (urls: string[], jobId?: string) => {
    const newImages: StoredImage[] = urls.map(url => ({
      url,
      timestamp: Date.now(),
      pageType,
      jobId
    }));

    setState((prevState) => ({
      lastJobId: jobId,
      images: [...newImages, ...prevState.images.slice(0, 23)] // Keep max 24 images
    }));
  };

  const clearImages = () => {
    setState({ images: [] });
  };

  const getLatestImages = (count: number = 1): string[] => {
    return state.images.slice(0, count).map(img => img.url);
  };

  return {
    images: state.images,
    lastJobId: state.lastJobId,
    addImages,
    clearImages,
    getLatestImages
  };
}