backsplash-app/src/ui/hooks/useWallpaperGeneration.ts

An AI powered wallpaper app.

import { useState, useEffect, useCallback, useRef } from "react";
import { useLastWallpaper, useCurrentWallpaperMode } from "@/ui/hooks";
import { ServerChannels } from "@/ipc/channels/serverChannels";
import { WallpaperChannels, WallpaperStatusType } from "@/ipc/channels/wallpaperChannels";
import { WallpaperData, WallpaperServiceGenerationResponse } from "@/types/wallpaper";

// Define StatusUpdate locally or import from a shared types location if it exists beyond Recoil atom
export interface StatusUpdatePayload {
  type: WallpaperStatusType | string; // Allow string for custom/other status types from main
  message: string;
  data?: any;
  timestamp?: string;
}

// Define store key for wallpaper generation state
const WALLPAPER_GENERATION_STATE_KEY = "wallpaper-generation-state";
// Limit the number of status updates we keep in memory
const MAX_STATUS_UPDATES = 20;

interface WallpaperGenerationState {
  isGenerating: boolean;
  isUpscaling: boolean;
  error: string | null;
}

export function useWallpaperGeneration() {
  const [isGenerating, setIsGenerating] = useState<boolean>(false);
  const [isUpscaling, setIsUpscaling] = useState<boolean>(false);
  const [error, setError] = useState<string | null>(null);
  const [statusUpdates, setStatusUpdates] = useState<StatusUpdatePayload[]>([]);
  const [latestStatus, setLatestStatus] = useState<StatusUpdatePayload | null>(null);

  // Use refs to avoid useEffect dependency issues
  const isGeneratingRef = useRef(isGenerating);
  const isUpscalingRef = useRef(isUpscaling);
  const errorRef = useRef(error);

  // Keep refs in sync with state
  useEffect(() => {
    isGeneratingRef.current = isGenerating;
    isUpscalingRef.current = isUpscaling;
    errorRef.current = error;
  }, [isGenerating, isUpscaling, error]);

  const { setLastWallpaper } = useLastWallpaper();
  const { currentMode } = useCurrentWallpaperMode();

  // Load persisted state on component mount
  useEffect(() => {
    // Fetch initial state from electron-store
    window.electron.ipcRenderer
      .invoke("store:get", WALLPAPER_GENERATION_STATE_KEY)
      .then((state: WallpaperGenerationState | undefined) => {
        if (state) {
          setIsGenerating(state.isGenerating ?? false);
          setIsUpscaling(state.isUpscaling ?? false);
          setError(state.error ?? null);
        }
      })
      .catch((err) => console.error("Error fetching wallpaper generation state:", err));

    // Listen for updates from main process or other renderers
    const handleStoreUpdate = (_event: any, key: string, newValue: any) => {
      if (key === WALLPAPER_GENERATION_STATE_KEY && newValue) {
        try {
          // Try to parse if it's a string
          const state = typeof newValue === "string" ? JSON.parse(newValue) : newValue;
          setIsGenerating(state.isGenerating ?? isGeneratingRef.current);
          setIsUpscaling(state.isUpscaling ?? isUpscalingRef.current);
          setError(state.error ?? errorRef.current);
        } catch (err) {
          console.error("Error parsing wallpaper generation state:", err);
        }
      }
    };

    window.electron.ipcRenderer.on("store:updated", handleStoreUpdate);

    return () => {
      window.electron.ipcRenderer.removeListener("store:updated", handleStoreUpdate);
    };
  }, []);

  useEffect(() => {
    const handleStatusUpdate = (_event: any, update: StatusUpdatePayload) => {
      // Add timestamp if not present
      const newStatus = { ...update, timestamp: update.timestamp || new Date().toISOString() };

      // Update statusUpdates with a limit to prevent memory leaks
      setStatusUpdates((prev) => {
        const newUpdates = [...prev, newStatus];
        // Keep only the most recent MAX_STATUS_UPDATES
        return newUpdates.slice(-MAX_STATUS_UPDATES);
      });

      setLatestStatus(newStatus);

      switch (update.type) {
        case "generation-start":
          setIsGenerating(true);
          setError(null);
          break;
        case "generation-complete":
          setIsGenerating(false);
          if (!update.data?.success) {
            setError(update.data?.error || "Generation failed");
          }
          break;
        case "upscale-start":
          setIsUpscaling(true);
          setError(null);
          break;
        case "upscale-complete":
        case "upscale-error": // Handle upscale errors reported via status update
          setIsUpscaling(false);
          if (!update.data?.success) setError(update.data?.error || "Upscaling failed");
          break;
        default:
          break;
      }
    };
    window.electron.ipcRenderer.on(WallpaperChannels.WALLPAPER_STATUS_UPDATE, handleStatusUpdate);
    return () => {
      window.electron.ipcRenderer.removeListener(WallpaperChannels.WALLPAPER_STATUS_UPDATE, handleStatusUpdate);
    };
  }, []);

  const handleGenerate = useCallback(async () => {
    setIsGenerating(true);
    setError(null);

    try {
      // Ensure current mode is set in the store before generation
      if (!currentMode) {
        throw new Error("Current wallpaper mode is not selected.");
      }

      // GenerationService will use StoreService to get all necessary params
      const result = (await window.electron.ipcRenderer.invoke(
        ServerChannels.GENERATE_WALLPAPER,
      )) as WallpaperServiceGenerationResponse;

      if (!result || !result.success) {
        const errorDetail = (result as any)?.error || "Unknown error during generation";
        throw new Error(`Wallpaper generation failed: ${errorDetail}`);
      }

      // Update the last wallpaper with the generated data
      const wallpaperData: WallpaperData = {
        objectKey: result.objectKey,
        upscaledImageUrl: null,
        originalImageUrl: result.imageUrl,
        style: result.metadata?.style || "Unknown",
        generatedAt: new Date().toISOString(),
        metadata: result.metadata,
      };
      setLastWallpaper(wallpaperData);
    } catch (e: any) {
      console.error("Error in handleGenerate:", e);
      setError(e.message || "Failed to generate wallpaper.");
      setIsGenerating(false);
    }
  }, [currentMode, setLastWallpaper]);

  const handleSetWallpaper = useCallback(async (wallpaperToSet: WallpaperData) => {
    const dataToUse = wallpaperToSet;
    if (!dataToUse) {
      setError("No wallpaper data to set.");
      return;
    }

    setIsUpscaling(true);
    setError(null);

    try {
      const ipcResult = (await window.electron.ipcRenderer.invoke(
        WallpaperChannels.SET_WALLPAPER,
        dataToUse.objectKey,
      )) as { success: boolean; isUpscaled?: boolean; error?: string };

      if (!ipcResult.success) {
        throw new Error(ipcResult.error || "Failed to set wallpaper.");
      }
      // isUpscaling will be set to false by the 'upscale-complete' or 'upscale-error' status update listener.
    } catch (e: any) {
      console.error("Error setting wallpaper:", e);
      setIsUpscaling(false);
      setError(e.message || "Failed to set wallpaper.");
    }
  }, []);

  const handleUpscaleWallpaper = useCallback(
    async (wallpaperToUpscale: WallpaperData) => {
      const dataToUse = wallpaperToUpscale;
      if (!dataToUse) {
        setError("No wallpaper data to upscale.");
        return;
      }
      if (!dataToUse.objectKey) {
        setError("Wallpaper data has no objectKey to upscale.");
        return;
      }

      setIsUpscaling(true);
      setError(null);
      try {
        const result = (await window.electron.ipcRenderer.invoke(
          WallpaperChannels.UPSCALE_WALLPAPER,
          dataToUse.objectKey,
        )) as { success: boolean; error?: string; imageUrl?: string; objectKey?: string };
        if (!result.success) {
          throw new Error(result.error || "Failed to upscale wallpaper.");
        }
        // If successful, update localWallpaperData with new upscaledImageUrl if available from result
        // The primary state update (isUpscaling=false) comes from WALLPAPER_STATUS_UPDATE listener.
        if (result.imageUrl && result.objectKey === dataToUse.objectKey) {
          const updatedWallpaper = { ...dataToUse, upscaledImageUrl: result.imageUrl };
          // Persist this change to the store (lastWallpaper)
          setLastWallpaper(updatedWallpaper);
        }
      } catch (e: any) {
        console.error("Error upscaling wallpaper:", e);
        setIsUpscaling(false);
        setError(e.message || "Failed to upscale wallpaper.");
      }
    },
    [setLastWallpaper],
  );

  const resetError = useCallback(() => setError(null), []);
  const clearStatusUpdates = useCallback(() => setStatusUpdates([]), []);

  return {
    isGenerating,
    isUpscaling,
    error,
    statusUpdates,
    latestStatus,
    handleGenerate,
    handleSetWallpaper,
    handleUpscaleWallpaper,
    resetError,
    clearStatusUpdates,
  };
}