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

import { useState, useEffect } 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;
}