backsplash-app/src/ipc/services/generationService.ts

An AI powered wallpaper app.

import log from "@/logger";
import { getApiService } from "@/ipc/services/apiServiceInstance";
import { getWallpaperImageUrl } from "@/ipc/apiService";
import { FilesystemStorageService } from "@/ipc/services/filesystemStorageService";
import { StoreService } from "@/ipc/services/storeService";
import {
  WallpaperGenerationResponse,
  WallpaperLocationData,
  WallpaperServiceGenerationResponse,
  WallpaperServiceGenerationFailure,
  WallpaperServiceUpscaleResponse,
} from "@/types/wallpaper";
import { WallpaperMode } from "@/types/wallpaperModeTypes";
import { sendStatusUpdate } from "@/ipc/wallpaperIPC";

// Define the state interface (same as in the hook)
interface WallpaperGenerationState {
  isGenerating: boolean;
  isUpscaling: boolean;
  error: string | null;
}

/**
 * Generation Service
 *
 * Handles fetching, generating, and upscaling wallpapers.
 */
export class GenerationService {
  private fileSystem: FilesystemStorageService;
  private storeService: StoreService;
  private getStyleMap: () => Record<string, string>;
  private setStyleMap: (styleMap: Record<string, string>) => void;
  private apiService: ReturnType<typeof getApiService>;

  constructor(
    storeService: StoreService,
    fileSystem: FilesystemStorageService,
    getStyleMap: () => Record<string, string>,
    setStyleMap: (styleMap: Record<string, string>) => void,
  ) {
    this.storeService = storeService;
    this.fileSystem = fileSystem;
    this.getStyleMap = getStyleMap;
    this.setStyleMap = setStyleMap;
    this.apiService = getApiService(storeService);

    // Reset generation state on initialization
    const initialGenerationState: WallpaperGenerationState = {
      isGenerating: false,
      isUpscaling: false,
      error: null,
    };
    this.storeService.set("wallpaper-generation-state", initialGenerationState);

    log.info(`[GenerationService] Initialized and generation state reset.`);
  }

  /**
   * Updates the wallpaper generation state in the store
   */
  private updateGenerationState(updates: Partial<WallpaperGenerationState>): void {
    try {
      // Get current state from store or create default
      const currentState = (this.storeService.get<"wallpaper-generation-state">(
        "wallpaper-generation-state",
      ) as WallpaperGenerationState) || {
        isGenerating: false,
        isUpscaling: false,
        error: null,
      };

      // Create a new state object with only the properties we want to track
      const newState: WallpaperGenerationState = {
        isGenerating: updates.isGenerating ?? currentState.isGenerating,
        isUpscaling: updates.isUpscaling ?? currentState.isUpscaling,
        error: updates.error ?? currentState.error,
      };

      // Persist to store
      this.storeService.set("wallpaper-generation-state", newState);
      log.debug(`[GenerationService] Updated generation state: ${JSON.stringify(updates)}`);
    } catch (err) {
      log.error(`[GenerationService] Error updating generation state:`, err);
    }
  }

  /**
   * Downloads an image from a given URL and saves it locally.
   * @param imageUrl URL of the image to download.
   * @param filename Desired filename for the saved image.
   * @returns The local path where the image was saved.
   */
  private async downloadImage(imageUrl: string, filename: string): Promise<string> {
    log.info(`[GenerationService] Downloading image from ${imageUrl}`);

    if (!imageUrl) {
      const error = new Error("Empty image URL provided to downloadImage");
      log.error(`[GenerationService] ${error.message}`);
      throw error;
    }

    try {
      // Before downloading, verify URL format
      const urlObj = new URL(imageUrl);
      log.info(`[GenerationService] Valid URL: ${urlObj.toString()}`);

      const imagePath = await this.fileSystem.downloadFile(imageUrl, filename);
      log.info(`[GenerationService] Image downloaded successfully to ${imagePath}`);
      return imagePath;
    } catch (error) {
      if (error instanceof Error && error.message.includes("Invalid URL")) {
        log.error(`[GenerationService] Invalid image URL format: ${imageUrl}`);
      } else if (error instanceof Error && error.message.includes("404")) {
        log.error(`[GenerationService] Image not found at URL (404): ${imageUrl}`);
      } else {
        log.error(`[GenerationService] Error downloading image:`, error);
      }
      throw error;
    }
  }

  /**
   * Generate a new wallpaper. Fetches mode and parameters from StoreService.
   */
  async generateWallpaper(): Promise<WallpaperServiceGenerationResponse | WallpaperServiceGenerationFailure> {
    // Update generation state in store
    this.updateGenerationState({ isGenerating: true, error: null });

    const currentMode = this.storeService.getCurrentMode();
    log.info(`[GenerationService] Generating wallpaper for mode from store: ${currentMode}`);

    if (!currentMode) {
      const errorMsg = "Wallpaper generation failed: Current mode not set in store.";
      log.error(`[GenerationService] ${errorMsg}`);

      // Update generation state in store
      this.updateGenerationState({ isGenerating: false, error: errorMsg });

      sendStatusUpdate({
        type: "generation-complete",
        message: errorMsg,
        data: { success: false, error: errorMsg },
      });
      return { success: false, error: errorMsg };
    }

    sendStatusUpdate({
      type: "generation-start",
      message: `Starting wallpaper generation for mode: ${currentMode}`,
    });

    try {
      let styleParam = "";
      let customStyleParam: string | undefined;
      let locationDataParam: WallpaperLocationData | undefined;
      let modeParam: WallpaperMode | string | undefined = currentMode; // API mode param
      let actualStyleForMap: string = currentMode; // Style to save in styleMap

      switch (currentMode) {
        case "RandomStyleSelection": {
          const activeStyles = this.storeService.getActiveStyles();
          if (!activeStyles || activeStyles.length === 0) {
            log.warn(
              "[GenerationService] No active styles for RandomStyleSelection mode. Defaulting to 'Moody Gradient Spheres' style.",
            );
            styleParam = "Moody Gradient Spheres"; // Using a valid style from the server's list
            actualStyleForMap = styleParam;
            modeParam = undefined; // API uses the style directly
          } else {
            // Check if the selected style is in the valid list
            // If not, use a safe default
            styleParam = activeStyles[Math.floor(Math.random() * activeStyles.length)];
            // List of valid styles based on server error message
            const validStyles = [
              "Hyperrealistic Mountains",
              "Dramatic Coastlines",
              "Ancient Forests",
              "Surreal Skies",
              "Multiple Moons",
              "Aurora Phenomena",
              "Underwater Coral Cities",
              "Crystal Caves",
              "Magma Rivers",
              "Nebula Gardens",
              "Black Hole Event Horizons",
              "Stellar Nurseries",
              "Galactic Core",
              "Interstellar Clouds",
              "Ringed Gas Giants",
              "Dying Stars",
              "Cosmic Jellyfish",
              "Solar Prominences",
              "Cyberpunk Cityscape",
              "Neon-Soaked Megalopolis",
              "Gigantic Space Stations",
              "Orbital Habitats",
              "Retro-Futuristic Highways",
              "Alien Technology",
              "Quantum Realms",
              "Post-Apocalyptic Wasteland",
              "Biopunk Evolution",
              "Floating Islands",
              "Elven Forest Cities",
              "Dragon Realms",
              "Ancient Castles",
              "Mythical Beasts",
              "Celestial Kingdoms",
              "Magical Portal Worlds",
              "Crystal Fortresses",
              "Arcane Libraries",
              "Geometric Symmetry",
              "Bold Color Fields",
              "Moody Gradient Spheres",
              "Fractal Dimensions",
              "AI-Generated Mandalas",
              "Digital Void Art",
              "Prismatic Reflections",
              "Particle Systems",
              "Liquid Dreams",
              "Brutalist Megastructures",
              "Floating Island Homes",
              "Zen Minimalist Interiors",
              "Escher Dreamscapes",
              "Hollow Earth Cities",
              "Living Buildings",
              "Impossible Architecture",
              "Crystalline Towers",
              "Bio-Organic Structures",
              "Cyberpunk Portraits",
              "AI-Generated Deities",
              "Futuristic Knights",
              "Mage Warriors",
              "Tribal Sci-Fi Hunters",
              "Digital Soul Entities",
              "Technological Shamans",
              "Cosmic Travelers",
              "Hybrid Beings",
              "Void Landscapes",
              "Glowing Abyss",
              "Eerie Haunted Forests",
              "Dystopian Rain Cities",
              "Abstract Horror",
              "Shadow Realms",
              "Cosmic Horror",
              "Neon Noir",
              "Dark Dimension",
              "Synthwave Sunsets",
              "CRT Glitch Dreamscapes",
              "1980s Tech Nostalgia",
              "Neo-Tokyo Nights",
              "Retrowave Grids",
              "VHS Memories",
              "Digital Sunset",
              "Outrun Aesthetics",
              "20XX Vibes",
              "Surrealist Dreamscapes",
              "Digital Impressionism",
              "Hyper Detailed Realism",
              "Oil Painting Universe",
              "Watercolor Worlds",
              "AI Cubism",
              "Neural Style Transfer",
              "Generative Patterns",
              "Stable Diffusion Dreams",
              "Organic Machinery",
              "Techno-Organic Fusion",
              "Living Technology",
              "Neural Circuit Gardens",
              "Digital Organisms",
              "Synthetic Biology",
              "Machine Ecosystems",
              "Sentient Hardware",
              "Cybernetic Nature",
              "Liminal Spaces",
              "Memory Fragments",
              "Nostalgia Distortion",
              "Uncanny Familiarity",
              "Altered Reality",
              "Subconscious Landscapes",
              "Dream Logic Architecture",
              "Childhood Memories",
              "Fever Dream",
              "Ultra-Detailed Closeups",
              "Hyperreal Urban Scenes",
              "Perfect Reflections",
              "Studio Lighting Perfection",
              "Impossible Photography",
              "Macro Universe",
              "Microscopic Worlds",
              "Golden Hour Forever",
              "Texture Studies",
              "Neon Sacred Geometry",
              "Hyperdimensional Patterns",
              "Liquid Light Show",
              "Fractal Recursion",
              "Kaleidoscopic Visions",
              "Vibrating Colors",
              "Consciousness Expansion",
              "Quantum Visuals",
              "DMT Entities",
              "Fire Vortex",
              "Water Dimension",
              "Earth Consciousness",
              "Air Kingdom",
              "Lightning Networks",
              "Ice Architecture",
              "Plasma Forms",
              "Metal Living",
              "Ethereal Elements",
              "Light Refraction",
              "Digital Holograms",
              "Iridescent Surfaces",
              "Transparent Data",
              "3D Light Sculptures",
              "Rainbow Interfaces",
              "Projected Realities",
              "Prism Fields",
              "Holographic Universe",
              "Gene-Spliced Organisms",
              "Organic Computing",
              "Biotech Revolution",
              "Living Architecture",
              "Grown Nanomachines",
              "Genetic Art",
              "Engineered Ecosystems",
              "Fungal Networks",
              "Pulsing Biomass",
              "Custom",
              "Mixed",
              "Location",
              "Dall-E Style",
              "Midjourney Aesthetic",
              "Stable Diffusion XL",
              "AI Dream",
              "Photographic",
              "Cinematic",
            ];

            // Validate the selected style against the list of valid styles
            if (!validStyles.includes(styleParam)) {
              log.warn(
                `[GenerationService] Selected style '${styleParam}' is not in the server's valid list. Using 'Moody Gradient Spheres' instead.`,
              );
              styleParam = "Moody Gradient Spheres"; // Fallback to a known valid style
            }

            actualStyleForMap = styleParam; // The selected or fallback style
            modeParam = undefined; // Use style directly for API
            log.debug(`[GenerationService] Using style: ${styleParam}`);
          }
          break;
        }

        case "CustomStyle": {
          customStyleParam = this.storeService.getModeParams(currentMode)?.customStyle;
          if (!customStyleParam) {
            const errorMsg = "Generation failed: Custom style not set in store for CustomStyle mode.";
            log.error(`[GenerationService] ${errorMsg}`);
            throw new Error(errorMsg);
          }
          styleParam = "Custom"; // API expects "Custom" as style when custom_style is provided
          actualStyleForMap = "Custom";
          // modeParam is already "CustomStyle", which is fine if API handles it, or set to undefined if styleParam="Custom" is enough
          log.debug(`[GenerationService] Using custom style: ${customStyleParam}`);
          break;
        }

        case "Location": {
          const locationParams = this.storeService.getModeParams(currentMode);
          if (!locationParams?.locationData?.city) {
            const errorMsg = "Generation failed: Location data (city) not set in store for Location mode.";
            log.error(`[GenerationService] ${errorMsg}`);
            throw new Error(errorMsg);
          }
          locationDataParam = {
            city: locationParams.locationData.city,
            state: locationParams.locationData.state || "",
            country: locationParams.locationData.country || "",
            local_time_str: locationParams.locationData.local_time_str || new Date().toLocaleString(),
          };
          styleParam = currentMode; // Set styleParam to "Location"
          modeParam = undefined; // Do not send a separate mode if style implies it
          actualStyleForMap = "Location";
          log.debug(`[GenerationService] Using location data:`, locationDataParam, `and style: ${styleParam}`);
          break;
        }

        // Add cases for other modes like "Mood", "Weather", fetching their specific params from store
        case "Mood": {
          const moodParams = this.storeService.getModeParams(currentMode);
          if (!moodParams?.mood) {
            throw new Error("Mood not set for Mood mode.");
          }

          // Map mood to appropriate style from the valid list
          const moodStyleMap: Record<string, string> = {
            calm: "Zen Minimalist Interiors",
            energetic: "Neon Sacred Geometry",
            peaceful: "Watercolor Worlds",
            mysterious: "Shadow Realms",
            happy: "Synthwave Sunsets",
            sad: "Moody Gradient Spheres",
            angry: "Fire Vortex",
            dreamy: "Surrealist Dreamscapes",
            nostalgic: "VHS Memories",
            futuristic: "Cyberpunk Cityscape",
            natural: "Ancient Forests",
            urban: "Neo-Tokyo Nights",
            cosmic: "Stellar Nurseries",
            abstract: "Abstract Horror",
            magical: "Magical Portal Worlds",
          };

          // Get the mapped style or default to a known valid style
          styleParam = moodStyleMap[moodParams.mood.toLowerCase()] || "Moody Gradient Spheres";
          actualStyleForMap = styleParam;
          modeParam = undefined; // API uses the style directly
          log.debug(`[GenerationService] Mood mode using style: ${styleParam} for mood: ${moodParams.mood}`);
          break;
        }

        default:
          log.debug(`[GenerationService] Using direct style/mode from store: ${currentMode}`);
          // For any other modes, we'll use a safe default style
          // Only use the currentMode directly if it's one of our special modes
          if (currentMode === "Location" || currentMode === "Custom" || currentMode === "Mixed") {
            styleParam = currentMode;
          } else {
            // Default to a known good style
            styleParam = "Moody Gradient Spheres";
            log.warn(`[GenerationService] Unknown mode '${currentMode}', using default style '${styleParam}'`);
          }
          actualStyleForMap = currentMode; // Keep the original mode in our records
          modeParam = undefined;
          break;
      }

      // Now undefined for Location mode

      // Create request data object based on params
      const requestData = {
        style: styleParam,
        custom_style: customStyleParam,
        location_data: locationDataParam,
        mode: modeParam,
      };

      log.info(`[GenerationService] Generating with params:`, requestData);

      const result: WallpaperGenerationResponse = await this.apiService.wallpapers.generate(requestData);
      log.info(`[GenerationService] Generation API successful. Key: ${result.objectKey}`);

      this.saveStyle(result.objectKey, actualStyleForMap);

      const imageUrl = getWallpaperImageUrl(result.objectKey);
      const metadata = result.metadata || {
        style: actualStyleForMap,
        created_at: new Date().toISOString(),
        prompt: "",
        original_prompt: "",
      };

      this.storeService.setLastWallpaper({
        objectKey: result.objectKey,
        upscaledImageUrl: null,
        originalImageUrl: imageUrl,
        style: actualStyleForMap,
        generatedAt: new Date().toISOString(),
      });

      // Update generation state in store - generation complete
      this.updateGenerationState({ isGenerating: false, error: null });

      sendStatusUpdate({
        type: "generation-complete",
        message: "Wallpaper generation completed successfully",
        data: { success: true, objectKey: result.objectKey },
      });

      return {
        success: true,
        imageUrl,
        objectKey: result.objectKey,
        metadata,
      };
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      log.error(`[GenerationService] Error in generateWallpaper: ${errorMessage}`);

      // Update generation state in store - generation failed
      this.updateGenerationState({ isGenerating: false, error: errorMessage });

      sendStatusUpdate({
        type: "generation-complete",
        message: `Wallpaper generation failed: ${errorMessage}`,
        data: { success: false, error: errorMessage },
      });
      return { success: false, error: errorMessage };
    }
  }

  /**
   * Upscale an existing wallpaper
   */
  async upscaleWallpaper(objectKey: string): Promise<WallpaperServiceUpscaleResponse> {
    // Update upscaling state in store
    this.updateGenerationState({ isUpscaling: true, error: null });

    log.info(`[GenerationService] Upscaling wallpaper: ${objectKey}`);

    // Send upscale start status update
    sendStatusUpdate({
      type: "upscale-start",
      message: `Starting wallpaper upscaling for ${objectKey}`,
    });

    try {
      // Get the current model provider from store or default to 'replicate'
      const modelProvider = (this.storeService.get("model-provider") as string) || "replicate";

      // Get the license info
      const licenseKey = (this.storeService.get("license-key") as string) || "";
      const licensePlan = (this.storeService.get("license-plan") as string) || "free";
      const clientId = (this.storeService.get("client-id") as string) || "";

      // Call the API to upscale the image
      const response = await this.apiService.wallpapers.upscale({
        object_key: objectKey,
        model_provider: modelProvider,
        license_key: licenseKey,
        license_plan: licensePlan,
        client_id: clientId,
      });

      log.info(`[GenerationService] Successfully called upscale API for ${objectKey}`);

      // Check if we have an image URL to download
      if (response.imageUrl) {
        try {
          // Get upscaled filename based on the object key
          const { upscaledFilename } = this.fileSystem.getWallpaperUpscaledInfo(objectKey);

          log.info(`[GenerationService] Downloading upscaled image to: ${upscaledFilename}`);
          log.info(`[GenerationService] Using image URL: ${response.imageUrl}`);

          const downloadedPath = await this.downloadImage(response.imageUrl, upscaledFilename);

          // Verify the download succeeded
          if (this.fileSystem.fileExists(downloadedPath)) {
            log.info(`[GenerationService] Successfully downloaded upscaled image to: ${downloadedPath}`);

            // Get original image info to preserve style and other metadata using API
            const originalImageUrl = getWallpaperImageUrl(objectKey);
            const upscaledImageUrl = getWallpaperImageUrl(downloadedPath);

            log.info(`[GenerationService] Original image URL: ${originalImageUrl}`);
            log.info(`[GenerationService] Upscaled image URL: ${upscaledImageUrl}`);

            // Fetch the metadata from the API to ensure we have accurate style info
            log.info(`[GenerationService] Fetching metadata for: ${objectKey}`);
            const metadata = await this.apiService.wallpapers.getMetadata(objectKey);
            log.info(`[GenerationService] Retrieved metadata: ${JSON.stringify(metadata, null, 2)}`);

            // Use saved style if available, otherwise use the metadata or inferred style
            const styleInfo = this.getSavedStyle(objectKey) || metadata.style || "Unknown";
            log.info(
              `[GenerationService] Using style: ${styleInfo} (from ${this.getSavedStyle(objectKey) ? "saved map" : metadata.inferred ? "inferred" : "metadata"})`,
            );

            // Update upscaling state in store - success
            this.updateGenerationState({ isUpscaling: false, error: null });

            // Preserve style information in status update
            sendStatusUpdate({
              type: "upscale-complete",
              message: `Completed wallpaper upscaling for ${objectKey}`,
              data: {
                success: true,
                isUpscaled: true,
                objectKey: objectKey,
                style: styleInfo,
                originalImageUrl,
                upscaledImageUrl,
              },
            });

            return {
              success: true,
              objectKey: objectKey,
              imageUrl: upscaledImageUrl,
            };
          } else {
            log.error(`[GenerationService] Failed to verify downloaded file exists: ${downloadedPath}`);

            // Update upscaling state in store - failed
            const errorMsg = "Failed to verify downloaded file exists";
            this.updateGenerationState({ isUpscaling: false, error: errorMsg });

            return {
              success: false,
              error: errorMsg,
              objectKey,
            };
          }
        } catch (downloadError) {
          log.error(`[GenerationService] Error downloading upscaled image:`, downloadError);

          // Update upscaling state in store - download failed
          const errorMsg = downloadError instanceof Error ? downloadError.message : "Failed to download upscaled image";
          this.updateGenerationState({ isUpscaling: false, error: errorMsg });

          // Return failure with error
          return {
            success: false,
            error: errorMsg,
            objectKey,
          };
        }
      } else {
        log.warn(`[GenerationService] No image URL in upscale result for ${objectKey}`);

        // Update upscaling state in store - no URL
        const errorMsg = "No image URL in upscale result";
        this.updateGenerationState({ isUpscaling: false, error: errorMsg });

        return {
          success: false,
          error: errorMsg,
          objectKey,
        };
      }
    } catch (error) {
      log.error(`[GenerationService] Error upscaling wallpaper ${objectKey}:`, error);

      // Update upscaling state in store - general error
      const errorMsg = error instanceof Error ? error.message : "Upscaling failed";
      this.updateGenerationState({ isUpscaling: false, error: errorMsg });

      // Send upscale failure status update
      sendStatusUpdate({
        type: "upscale-complete",
        message: `Wallpaper upscaling failed for ${objectKey}`,
        data: {
          success: false,
          error: errorMsg,
        },
      });

      return {
        success: false,
        error: errorMsg,
        objectKey,
      };
    }
  }

  private getSavedStyle(objectKey: string): string | undefined {
    const styleMap = this.getStyleMap();
    return styleMap[objectKey];
  }

  private saveStyle(objectKey: string, style: string): void {
    const styleMap = this.getStyleMap();
    styleMap[objectKey] = style;
    this.setStyleMap(styleMap);
    log.debug(`[GenerationService] Saved style "${style}" for ${objectKey} in style map`);
  }
}