backsplash-app/src/ipc/services/wallpaperService.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 {
  WallpaperServiceGenerationResponse,
  WallpaperServiceGenerationFailure,
  WallpaperServiceUpscaleResponse,
  WallpaperHistoryEntry,
  WallpaperMetadata,
} from "@/types/wallpaper";
import { BrowserWindow } from "electron";
import { WallpaperUpdates } from "../channels/wallpaperChannels";
import { GenerationService } from "./generationService";
import { Worker } from "worker_threads";
import path from "path";

// Constants for performance optimization
const BATCH_OPERATION_TIMEOUT = 30000; // 30 seconds
const MAX_RETRIES = 3;
const RETRY_DELAY = 1000;
const WORKER_CLEANUP_TIMEOUT = 5000; // 5 seconds

/**
 * Wallpaper Service
 *
 * Handles fetching, generating, upscaling, caching, and setting wallpapers.
 */
export class WallpaperService {
  private localCacheDir: string;
  private fileSystem: FilesystemStorageService;
  private storeService: StoreService;
  private readonly styleMapKey = "wallpaper-style-map";
  private readonly generationService: GenerationService;
  private worker: Worker | null = null;
  private messageQueue: Array<{ message: any; resolve: (value: any) => void; reject: (error: any) => void }> = [];
  private isProcessingQueue = false;
  private workerCleanupTimeout: NodeJS.Timeout | null = null;
  private apiService: ReturnType<typeof getApiService>;

  constructor(storeService: StoreService, fileSystem: FilesystemStorageService) {
    this.fileSystem = fileSystem;
    this.localCacheDir = this.fileSystem.getStorageDir();
    this.storeService = storeService;
    this.apiService = getApiService(storeService);
    this.generationService = new GenerationService(
      this.storeService,
      this.fileSystem,
      this.getStyleMap.bind(this),
      this.setStyleMap.bind(this),
    );

    log.info(`[WallpaperService] Cache directory: ${this.localCacheDir}`);
    log.info("[WallpaperService] GenerationService instantiated by WallpaperService.");
  }

  private initializeWorker() {
    if (this.worker) return;

    try {
      const workerPath = path.join(__dirname, "wallpaperWorker.js");
      log.info(`[WallpaperService] Attempting to initialize worker at: ${workerPath}`);

      this.worker = new Worker(workerPath);
      log.info(`[WallpaperService] Worker initialized successfully`);

      // Add error boundaries to prevent worker crashes
      this.worker.on("error", (error) => {
        log.error("[WallpaperService] Worker error:", error);
        const { reject } = this.messageQueue.shift() || {};
        if (reject) {
          reject(error);
        }
        this.cleanupWorker();
        this.processNextMessage();
      });

      this.worker.on("exit", (code) => {
        if (code !== 0) {
          log.error(`[WallpaperService] Worker stopped with exit code ${code}`);
        }
        this.cleanupWorker();
        this.processNextMessage();
      });

      this.worker.on("message", (response) => {
        const { resolve } = this.messageQueue.shift() || {};
        if (resolve) {
          resolve(response);
        }
        this.processNextMessage();
      });
    } catch (error) {
      log.error("[WallpaperService] Failed to initialize worker:", error);
      this.worker = null;
    }
  }

  private cleanupWorker() {
    if (this.worker) {
      this.worker.terminate();
      this.worker = null;
    }
    if (this.workerCleanupTimeout) {
      clearTimeout(this.workerCleanupTimeout);
      this.workerCleanupTimeout = null;
    }
  }

  private scheduleWorkerCleanup() {
    if (this.workerCleanupTimeout) {
      clearTimeout(this.workerCleanupTimeout);
    }
    this.workerCleanupTimeout = setTimeout(() => {
      if (this.messageQueue.length === 0) {
        this.cleanupWorker();
      }
    }, WORKER_CLEANUP_TIMEOUT);
  }

  private async processNextMessage() {
    if (this.isProcessingQueue || this.messageQueue.length === 0) return;

    this.isProcessingQueue = true;
    try {
      const { message, reject } = this.messageQueue[0];

      // Initialize worker asynchronously
      await new Promise<void>((resolveInit) => {
        this.initializeWorker();
        resolveInit();
      });

      if (this.worker) {
        this.worker.postMessage(message);
      } else {
        reject(new Error("Worker initialization failed"));
        this.messageQueue.shift();
      }
    } catch (error) {
      const { reject } = this.messageQueue.shift() || {};
      if (reject) {
        reject(error);
      }
    } finally {
      this.isProcessingQueue = false;
      this.scheduleWorkerCleanup();
    }
  }

  private async sendMessageToWorker(message: any): Promise<any> {
    return new Promise((resolve, reject) => {
      this.messageQueue.push({ message, resolve, reject });
      this.processNextMessage();
    });
  }

  private async batchFileSystemOperations<T>(operations: Array<() => Promise<T>>): Promise<T[]> {
    const timeoutPromise = new Promise<never>((_, reject) => {
      setTimeout(() => reject(new Error("Batch operation timeout")), BATCH_OPERATION_TIMEOUT);
    });

    try {
      // Reduce chunk size for better responsiveness
      const chunkSize = 3;
      const results: T[] = [];

      // Process chunks with a small delay between them to prevent UI blocking
      for (let i = 0; i < operations.length; i += chunkSize) {
        const chunk = operations.slice(i, i + chunkSize);
        const chunkResults = await Promise.race([Promise.all(chunk.map((op) => op())), timeoutPromise]);
        results.push(...chunkResults);

        // Add a small delay between chunks to allow UI updates
        if (i + chunkSize < operations.length) {
          await new Promise((resolve) => setTimeout(resolve, 50));
        }
      }

      return results;
    } catch (error) {
      log.error("[WallpaperService] Batch operation failed:", error);
      throw error;
    }
  }

  /**
   * Send a status update to all renderer processes (used internally or by GenerationService via constructor).
   * This method might be redundant if sendStatusUpdateFunc from wallpaperIPC is used directly by GenerationService.
   * However, keeping it here if WallpaperService itself needs to send distinct updates.
   */
  private sendInternalStatusUpdate(status: { type: string; message: string; data?: any }): void {
    const windows = BrowserWindow.getAllWindows();
    windows.forEach((window) => {
      if (!window.isDestroyed()) {
        // Assuming a generic channel or define a specific one for these internal updates
        window.webContents.send("wallpaper-service-status", status);
      }
    });
  }

  /**
   * Send a wallpaper update to all renderer processes
   */
  public sendWallpaperUpdate(update: {
    objectKey: string;
    originalImageUrl: string;
    upscaledImageUrl: string | null;
    setAsBackground?: boolean;
    style: string;
    metadata?: any;
  }): void {
    try {
      const windows = BrowserWindow.getAllWindows();
      windows.forEach((window) => {
        if (!window.isDestroyed()) {
          window.webContents.send(WallpaperUpdates.UPDATE, update);
        }
      });
      log.info(`[WallpaperService] Sent wallpaper update to renderer for key: ${update.objectKey}`);
    } catch (error) {
      log.error("[WallpaperService] Error sending wallpaper update:", error);
    }
  }

  /**
   * Set the wallpaper using the native system APIs.
   * Only uses upscaled versions - will fail if no upscaled version exists.
   * @param objectKey The unique identifier (e.g., S3 key/filename) for the wallpaper.
   * @returns Object with success status and upscaling information
   */
  async setWallpaper(objectKey: string): Promise<{ success: boolean; isUpscaled: boolean }> {
    log.info(`[WallpaperService] Request to set wallpaper for key: ${objectKey}`);

    let attempts = 0;
    while (attempts < MAX_RETRIES) {
      try {
        // Batch file system operations
        const [upscaledInfo, metadata, styleMap] = await this.batchFileSystemOperations([
          async () => this.fileSystem.getWallpaperUpscaledInfo(objectKey),
          async () => this.apiService.wallpapers.getMetadata(objectKey),
          async () => this.getStyleMap(),
        ]);

        const { upscaledPath } = upscaledInfo;
        const styleInfo = styleMap[objectKey] || metadata.style || "Unknown";
        let localPathToSet: string | null = null;

        // If we have a local upscaled version, use it
        if (upscaledPath) {
          log.info(`[WallpaperService] Using existing upscaled version locally: ${upscaledPath}`);
          localPathToSet = upscaledPath;
        } else {
          log.info(`[WallpaperService] No local upscaled version found for ${objectKey}. Performing upscale...`);

          let upscaleResult: { success: boolean; error?: string } = { success: false };

          // Try worker upscaling first if available
          if (this.worker) {
            try {
              log.info(`[WallpaperService] Attempting worker-based upscaling for ${objectKey}`);
              upscaleResult = await this.sendMessageToWorker({ type: "upscale", objectKey });
              log.info(`[WallpaperService] Worker upscaling result:`, upscaleResult);
            } catch (workerError) {
              log.error(`[WallpaperService] Worker upscaling failed:`, workerError);
              upscaleResult = { success: false, error: `Worker error: ${workerError}` };
            }
          }

          // Fallback to synchronous upscaling if worker failed or is not available
          if (!upscaleResult.success) {
            try {
              log.info(`[WallpaperService] Falling back to synchronous upscaling for ${objectKey}`);
              upscaleResult = await this.generationService.upscaleWallpaper(objectKey);
              log.info(`[WallpaperService] Synchronous upscaling result:`, upscaleResult);
            } catch (syncError) {
              log.error(`[WallpaperService] Synchronous upscaling failed:`, syncError);
              upscaleResult = { success: false, error: `Sync error: ${syncError}` };
            }
          }

          if (upscaleResult.success) {
            const { upscaledPath: freshUpscaledPath } = await this.fileSystem.getWallpaperUpscaledInfo(objectKey);
            if (freshUpscaledPath) {
              localPathToSet = freshUpscaledPath;
              log.info(`[WallpaperService] Successfully upscaled to: ${freshUpscaledPath}`);
            } else {
              log.error(`[WallpaperService] Upscaling succeeded but no upscaled file found for ${objectKey}`);
            }
          } else {
            log.error(`[WallpaperService] All upscaling attempts failed for ${objectKey}:`, upscaleResult.error);
          }
        }

        if (!localPathToSet) {
          log.error(`[WallpaperService] Could not determine a valid wallpaper path to set for key: ${objectKey}.`);
          return { success: false, isUpscaled: false };
        }

        // Get the image URLs for frontend
        const [upscaledImageUrl, originalImageUrl] = await this.batchFileSystemOperations([
          async () => getWallpaperImageUrl(localPathToSet!),
          async () => getWallpaperImageUrl(objectKey),
        ]);

        // Send a wallpaper update to the frontend
        this.sendWallpaperUpdate({
          objectKey,
          originalImageUrl,
          upscaledImageUrl,
          setAsBackground: true,
          style: styleInfo,
        });

        this.storeService.setLastWallpaper({
          objectKey,
          originalImageUrl,
          upscaledImageUrl,
          style: styleInfo,
          generatedAt: new Date().toISOString(),
        });

        // Set the wallpaper
        const wallpaperModule = await import("wallpaper");
        await wallpaperModule.setWallpaper(localPathToSet);
        log.info(`[WallpaperService] Successfully set wallpaper.`);

        // Store the current wallpaper path for scheduler enforcement
        this.storeService.set("current-wallpaper-path", localPathToSet);

        // Perform cache cleanup and history update in parallel
        await Promise.all([
          this.cleanupCache(localPathToSet),
          this.addToWallpaperHistory({
            key: objectKey,
            upscaledImageUrl,
            originalImageUrl,
            setAt: new Date().toISOString(),
            isUpscaled: true,
            metadata: {
              ...metadata,
              style: styleInfo,
            },
          }),
        ]);

        return { success: true, isUpscaled: true };
      } catch (error) {
        attempts++;
        if (attempts === MAX_RETRIES) {
          log.error(`[WallpaperService] Error setting wallpaper after ${MAX_RETRIES} attempts:`, error);
          return { success: false, isUpscaled: false };
        }
        log.warn(`[WallpaperService] Attempt ${attempts} failed, retrying in ${RETRY_DELAY}ms...`);
        await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY));
      }
    }
    return { success: false, isUpscaled: false };
  }

  /**
   * Generate a new wallpaper based on style, mode, etc. Delegates to GenerationService.
   * GenerationService will pull its parameters directly from the StoreService.
   */
  async generateWallpaper(): Promise<WallpaperServiceGenerationResponse | WallpaperServiceGenerationFailure> {
    return this.generationService.generateWallpaper();
  }

  /**
   * Generate a new wallpaper and immediately set it as the wallpaper.
   * @returns The result of the wallpaper generation.
   */
  async generateWallpaperAndImmediatelySet(): Promise<
    WallpaperServiceGenerationResponse | WallpaperServiceGenerationFailure
  > {
    const result = await this.generateWallpaper();
    if (result.success) {
      await this.setWallpaper(result.objectKey);
    }
    return result;
  }

  /**
   * Upscale an existing wallpaper. Delegates to GenerationService.
   */
  async upscaleWallpaper(objectKey: string): Promise<WallpaperServiceUpscaleResponse> {
    try {
      log.info(`[WallpaperService] Attempting to upscale wallpaper: ${objectKey}`);
      return await this.generationService.upscaleWallpaper(objectKey);
    } catch (error) {
      log.error(`[WallpaperService] Error upscaling wallpaper: ${objectKey}`, error);
      return {
        success: false,
        error: error instanceof Error ? error.message : "Unknown error occurred",
      };
    }
  }

  /**
   * Cleans up the wallpaper cache directory, keeping the N most recent files.
   * @param keepPath Optional path to explicitly keep, regardless of age.
   * @param keepCount Number of most recent files to keep (default: 5).
   */
  private async cleanupCache(keepPath?: string, keepCount = 5): Promise<void> {
    try {
      const pathsToKeep = keepPath ? [keepPath] : [];
      await this.fileSystem.cleanupOldFiles(pathsToKeep, keepCount);

      // Schedule next cleanup
      setTimeout(() => {
        this.cleanupCache().catch((error) => {
          log.error("[WallpaperService] Scheduled cleanup failed:", error);
        });
      }, 3600000); // Run cleanup every hour
    } catch (error) {
      log.error("[WallpaperService] Cache cleanup failed:", error);
    }
  }

  /**
   * Add a wallpaper to the history in electron-store
   */
  private async addToWallpaperHistory(wallpaper: {
    key: string;
    upscaledImageUrl: string;
    originalImageUrl: string;
    setAt: string;
    isUpscaled: boolean;
    metadata?: WallpaperMetadata;
  }): Promise<void> {
    log.debug("[WallpaperService] Adding to wallpaper history", wallpaper);
    try {
      const historyEntry: WallpaperHistoryEntry = {
        key: wallpaper.key,
        upscaledImageUrl: wallpaper.upscaledImageUrl,
        originalImageUrl: wallpaper.originalImageUrl,
        setAt: wallpaper.setAt,
        isUpscaled: wallpaper.isUpscaled,
        metadata: wallpaper.metadata,
      };

      // Batch store operations
      await Promise.all([
        this.storeService.addWallpaperHistoryEntry(historyEntry),
        this.updateStyleMap(wallpaper.key, wallpaper.metadata?.style),
      ]);

      // Notify renderer processes
      const windows = BrowserWindow.getAllWindows();
      const history = this.storeService.getWallpaperHistory();

      await Promise.all(
        windows.map((window) => {
          if (!window.isDestroyed() && window.webContents) {
            return window.webContents.send("wallpaper-history-updated", history);
          }
          return Promise.resolve();
        }),
      );
    } catch (error) {
      log.error("[WallpaperService] Failed to add to wallpaper history:", error);
      throw error;
    }
  }

  private async updateStyleMap(key: string, style?: string): Promise<void> {
    if (!style) return;

    const styleMap = this.getStyleMap();
    styleMap[key] = style;
    await this.setStyleMap(styleMap);
  }

  private getStyleMap(): Record<string, string> {
    const styleMap = this.storeService.get(this.styleMapKey);
    return styleMap && typeof styleMap === "object" && !Array.isArray(styleMap)
      ? (styleMap as Record<string, string>)
      : {};
  }

  private setStyleMap(newMap: Record<string, string>): void {
    this.storeService.set(this.styleMapKey, newMap as Record<string, string>);
  }

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