popiart-node-canvas/src/store/execution/nanoBananaExecutor.ts

/**
 * NanoBanana Executor
 *
 * Unified executor for nanoBanana (image generation) nodes.
 * Used by both executeWorkflow and regenerateNode.
 */

import type {
  NanoBananaNodeData,
  Resolution,
  SelectedModel,
} from "@/types";
import { calculateGenerationCost } from "@/utils/costCalculator";
import { buildGenerateHeaders } from "@/store/utils/buildApiHeaders";
import {
  createDefaultImageModel,
  createNewApiWGDefaultImageModel,
  createNewApiWGDefaultImageFallbackModel,
  isNewApiWGDefaultImageModel,
  shouldPreferNewApiWGImageModel,
} from "@/store/utils/defaultImageModel";
import { pollGenerateTask } from "./pollTaskCompletion";
import { runWithFallback } from "./runWithFallback";
import type { NodeExecutionContext } from "./types";
import {
  isBrowserFileSystemPath,
  joinBrowserFileSystemPath,
  writeBrowserGenerationFile,
} from "@/utils/browserFileSystem";
import { isHttpMediaUrl } from "@/utils/mediaResolver";
import {
  cleanupTemporaryImages,
  uploadPopiserverGenerationInputsForGeneration,
  uploadTemporaryDynamicInputsForGeneration,
  uploadTemporaryImagesForGeneration,
} from "@/utils/temporaryImageUpload";
import {
  mergeConfiguredParametersOverride,
  positiveInteger,
  resolveImageGenerationCount,
} from "./generationParameters";
import { buildDefaultParametersFromSchema, mergeSchemaDefaults } from "@/utils/modelSchemaDefaults";
import { readImageGenerationConfig, removeFirstClassGenerationParameters } from "@/utils/generationConfig";

export interface NanoBananaOptions {
  /** When true, falls back to stored inputImages/inputPrompt if no connections provide them. */
  useStoredFallback?: boolean;
}

type GeneratedResultImage = {
  url: string;
  previewImg?: string;
};

const GEMINI_RESOLUTIONS: Resolution[] = ["512", "1K", "2K", "4K"];

function normalizeGeminiResolution(value: string): Resolution {
  return GEMINI_RESOLUTIONS.includes(value as Resolution) ? (value as Resolution) : "2K";
}

function isProviderOverloadError(error: unknown): boolean {
  const message = error instanceof Error ? error.message : String(error);
  return /provider is temporarily overloaded|system\s+disk\s+overloaded/i.test(message);
}

async function loadSavedGenerationImage(directoryPath: string, imageId: string): Promise<string | null> {
  const response = await fetch("/api/load-generation", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ directoryPath, imageId }),
  });
  const result = await response.json();
  return result.success && typeof result.image === "string" ? result.image : null;
}

export async function executeNanoBanana(
  ctx: NodeExecutionContext,
  options: NanoBananaOptions = {}
): Promise<void> {
  const {
    node,
    getConnectedInputs,
    updateNodeData,
    getFreshNode,
    getEdges,
    getNodes,
    signal,
    providerSettings,
    providerEnvStatus,
    addIncurredCost,
    addToGlobalHistory,
    generationsPath,
    saveDirectoryPath,
    trackSaveGeneration,
    appendOutputGalleryImage,
  } = ctx;

  const { useStoredFallback = false } = options;

  const { images: connectedImages, text: connectedText, dynamicInputs } = getConnectedInputs(node.id);
  const effectiveGenerationsPath = generationsPath
    ?? (saveDirectoryPath && isBrowserFileSystemPath(saveDirectoryPath)
      ? joinBrowserFileSystemPath(saveDirectoryPath, "generations")
      : null);

  // Get fresh node data from store
  const freshNode = getFreshNode(node.id);
  const nodeData = (freshNode?.data || node.data) as NanoBananaNodeData;
  const defaultImageConfigModel = createDefaultImageModel(
    nodeData.model,
    providerSettings,
    providerEnvStatus
  );
  const nodeConfig = readImageGenerationConfig(nodeData, defaultImageConfigModel);

  // Determine images and text (with optional fallback to stored values)
  let images: string[];
  let promptText: string | null;

  if (useStoredFallback) {
    const localImage = (nodeData as NanoBananaNodeData & { image?: string | null }).image;
    images = connectedImages.length > 0
      ? connectedImages
      : localImage
        ? [localImage]
        : nodeConfig.inputImages;
    promptText = connectedText ?? nodeConfig.prompt;
  } else {
    const localImage = (nodeData as NanoBananaNodeData & { image?: string | null }).image;
    images = connectedImages.length > 0 ? connectedImages : localImage ? [localImage] : [];
    // For dynamic inputs, check if we have at least a prompt
    const promptFromDynamic = Array.isArray(dynamicInputs.prompt)
      ? dynamicInputs.prompt[0]
      : dynamicInputs.prompt;
    promptText = connectedText || promptFromDynamic || nodeConfig.prompt || null;
  }

  // Defensive: ensure promptText is actually a string at runtime
  // (Guards against corrupted node data or race conditions in parallel execution)
  if (promptText !== null && typeof promptText !== 'string') {
    const raw: unknown = promptText;
    console.warn('[nanoBanana] promptText was not a string, coercing:', typeof raw, Array.isArray(raw) ? `<redacted array length=${raw.length}>` : '<redacted>');
    promptText = Array.isArray(raw) ? (raw as string[])[0] ?? null : null;
  }

  if (!promptText) {
    updateNodeData(node.id, {
      status: "error",
      error: "Missing text input",
    });
    throw new Error("Missing text input");
  }

  // Capture promptText as a definitely-non-null string for use inside the closure.
  const finalPrompt: string = promptText;

  updateNodeData(node.id, {
    inputImages: images,
    inputPrompt: finalPrompt,
    status: "loading",
    error: null,
  });

  // Inner runOnce: performs the actual fetch/process/history work for a given model.
  // Extracted so runWithFallback can invoke it twice (primary, then fallback) if needed.
  const runOnce = async (modelToUse: SelectedModel, parametersOverride?: Record<string, unknown>): Promise<void> => {
    const provider = modelToUse.provider;
    const headers = buildGenerateHeaders(provider, providerSettings);

    // Sanitize dynamicInputs: remove prompt since it's already sent as the top-level
    // `prompt` field in requestPayload. Keeping both can cause providers like Replicate
    // to prefer dynamicInputs.prompt over the authoritative top-level value.
    const sanitizedDynamicInputs = { ...dynamicInputs };
    delete sanitizedDynamicInputs.prompt;

    let temporaryImageIds: string[] = [];

    try {
      const preparedPopiserverInputs = provider === "popiserver"
        ? await uploadPopiserverGenerationInputsForGeneration(images, sanitizedDynamicInputs)
        : null;
      const preparedImages = preparedPopiserverInputs
        ? { images: preparedPopiserverInputs.images, cleanupIds: preparedPopiserverInputs.cleanupIds }
        : provider === "newapiwg"
          ? await uploadTemporaryImagesForGeneration(images)
          : { images, cleanupIds: [] };
      const preparedDynamicInputs = preparedPopiserverInputs
        ? { dynamicInputs: preparedPopiserverInputs.dynamicInputs, cleanupIds: preparedPopiserverInputs.cleanupIds }
        : provider === "newapiwg"
          ? await uploadTemporaryDynamicInputsForGeneration(sanitizedDynamicInputs)
          : { dynamicInputs: sanitizedDynamicInputs, cleanupIds: [] };
      temporaryImageIds = [
        ...preparedImages.cleanupIds,
        ...preparedDynamicInputs.cleanupIds,
      ];

      const baseParameters = mergeSchemaDefaults(
        buildDefaultParametersFromSchema(nodeData.parameterSchema ?? []),
        parametersOverride ?? nodeConfig.parameters
      );
      const requestedImageCount =
        positiveInteger(nodeConfig.imageCount) ?? resolveImageGenerationCount(baseParameters, undefined);
      const normalizedBaseParameters = removeFirstClassGenerationParameters(baseParameters);
      const configuredParameters = {
        aspectRatio: (parametersOverride?.aspectRatio as string | undefined) ?? nodeConfig.aspectRatio,
        aspect_ratio: (parametersOverride?.aspectRatio as string | undefined) ?? nodeConfig.aspectRatio,
        ratio: (parametersOverride?.aspectRatio as string | undefined) ?? nodeConfig.aspectRatio,
        resolution: (parametersOverride?.resolution as string | undefined) ?? nodeConfig.resolution,
        imageSize: (parametersOverride?.resolution as string | undefined) ?? nodeConfig.resolution,
        useGoogleSearch: (parametersOverride?.useGoogleSearch as boolean | undefined) ?? nodeData.useGoogleSearch,
        useImageSearch: (parametersOverride?.useImageSearch as boolean | undefined) ?? nodeData.useImageSearch,
      };
      const effectiveParameters = provider === "popiserver"
        ? {
            ...mergeConfiguredParametersOverride(normalizedBaseParameters, configuredParameters),
            batchSize: requestedImageCount,
          }
        : mergeConfiguredParametersOverride(normalizedBaseParameters, configuredParameters);
      const usedModelId = modelToUse.modelId || nodeData.model;
      const requestPayload = {
        images: preparedImages.images,
        prompt: finalPrompt,
        aspectRatio: (parametersOverride?.aspectRatio as string) ?? nodeConfig.aspectRatio,
        resolution: (parametersOverride?.resolution as string) ?? nodeConfig.resolution,
        model: usedModelId,
        useGoogleSearch: (parametersOverride?.useGoogleSearch as boolean) ?? nodeData.useGoogleSearch,
        useImageSearch: (parametersOverride?.useImageSearch as boolean) ?? nodeData.useImageSearch,
        selectedModel: modelToUse,
        parameters: effectiveParameters,
        dynamicInputs: preparedDynamicInputs.dynamicInputs,
      };

      // Final guard: assert that prompt is a string before sending to API
      if (typeof requestPayload.prompt !== 'string') {
        const errorMsg = `Internal error: prompt is ${typeof requestPayload.prompt}, expected string`;
        console.error('[nanoBanana]', errorMsg);
        updateNodeData(node.id, { status: 'error', error: errorMsg });
        throw new Error(errorMsg);
      }

      const usePopiserverBatch = provider === "popiserver" && requestedImageCount > 1;
      const requestCount = usePopiserverBatch ? 1 : requestedImageCount;
      const usedModelDisplayName = modelToUse.displayName || usedModelId;
      const batchTimestamp = Date.now();
      const generatedItems: Array<{
        id: string;
        image: string;
        previewImg?: string;
        timestamp: number;
      }> = [];
      let partialError: string | null = null;

      for (let requestIndex = 0; requestIndex < requestCount; requestIndex++) {
        const response = await fetch("/api/generate", {
          method: "POST",
          headers,
          body: JSON.stringify(requestPayload),
          ...(signal ? { signal } : {}),
        });

        if (!response.ok) {
          const errorText = await response.text();
          let errorMessage = `HTTP ${response.status}`;
          try {
            const errorJson = JSON.parse(errorText);
            errorMessage = errorJson.error || errorMessage;
          } catch {
            if (errorText) errorMessage += ` - ${errorText.substring(0, 200)}`;
          }

          if (generatedItems.length > 0) {
            partialError = errorMessage;
            break;
          }

          updateNodeData(node.id, {
            status: "error",
            error: errorMessage,
          });
          throw new Error(errorMessage);
        }

        let result = await response.json();

        // Handle polling response (long-running Kie tasks)
        if (result.polling) {
          result = await pollGenerateTask({
            taskId: result.taskId,
            provider: result.pollProvider,
            modelId: result.pollModelId,
            modelName: result.pollModelName,
            mediaType: result.pollMediaType,
            headers,
            signal,
          });
        }
        const resultImages: GeneratedResultImage[] = result.success && Array.isArray(result.images) && result.images.length > 0
          ? result.images
          : result.success && result.image
            ? [result.image]
            : [];

        if (resultImages.length > 0) {
          resultImages.forEach((image) => {
            const itemIndex = generatedItems.length;
            const timestamp = batchTimestamp + itemIndex;
            const imageId = requestedImageCount === 1
              ? `${timestamp}`
              : `${batchTimestamp}-${itemIndex + 1}`;
            generatedItems.push({
              id: imageId,
              image: image.url,
              previewImg: image.previewImg,
              timestamp,
            });
          });
          continue;
        }

        const errorMessage = result.error || "Generation failed";
        if (generatedItems.length > 0) {
          partialError = errorMessage;
          break;
        }

        updateNodeData(node.id, {
          status: "error",
          error: errorMessage,
        });
        throw new Error(errorMessage);
      }

      if (usePopiserverBatch && generatedItems.length > 0 && generatedItems.length < requestedImageCount) {
        partialError = `Expected ${requestedImageCount} images, got ${generatedItems.length}`;
      }

      if (generatedItems.length === 0) {
        const errorMessage = partialError || "Generation failed";
        updateNodeData(node.id, {
          status: "error",
          error: errorMessage,
        });
        throw new Error(errorMessage);
      }

      const newHistoryItems = generatedItems.map((item) => ({
        id: item.id,
        image: item.image,
        previewImg: item.previewImg,
        timestamp: item.timestamp,
        prompt: finalPrompt,
        aspectRatio: nodeConfig.aspectRatio,
        model: usedModelId,
        modelDisplayName: usedModelDisplayName,
        modelProvider: modelToUse.provider,
      }));
      const updatedHistory = [...newHistoryItems, ...(nodeData.imageHistory || [])].slice(0, 50);

      updateNodeData(node.id, {
        outputImage: generatedItems[0].image,
        previewImg: generatedItems[0].previewImg,
        outputImageRef: undefined,
        outputImageStorageStatus: isHttpMediaUrl(generatedItems[0].image) ? "remote-only" : "localized",
        status: "complete",
        error: partialError,
        imageHistory: updatedHistory,
        selectedHistoryIndex: 0,
        lastUsedModel: modelToUse,
      });

      const edges = getEdges();
      const nodes = getNodes();

      generatedItems.forEach((item) => {
        // Save to global history
        addToGlobalHistory({
          image: item.image,
          timestamp: item.timestamp,
          prompt: finalPrompt,
          aspectRatio: nodeConfig.aspectRatio,
          model: usedModelId,
          modelDisplayName: usedModelDisplayName,
          modelProvider: modelToUse.provider,
        });

        // Push each generated image to connected downstream outputGallery nodes.
        edges
          .filter((e) => e.source === node.id)
          .forEach((e) => {
            const target = nodes.find((n) => n.id === e.target);
            if (target?.type === "outputGallery") {
              appendOutputGalleryImage(target.id, item.image);
            }
          });

        // Track cost per successful image.
        if (modelToUse.provider === "fal" && modelToUse.pricing) {
          addIncurredCost(modelToUse.pricing.amount);
        } else if (modelToUse.provider === "gemini") {
          const generationCost = calculateGenerationCost(nodeData.model, normalizeGeminiResolution(nodeConfig.resolution));
          addIncurredCost(generationCost);
        }

        // Auto-save to generations folder if configured
        if (effectiveGenerationsPath && isBrowserFileSystemPath(effectiveGenerationsPath)) {
          const savePromise = writeBrowserGenerationFile(
            effectiveGenerationsPath,
            item.id,
            item.image
          )
            .then((saveResult) => {
              const currentNode = getNodes().find((n) => n.id === node.id);
              if (!currentNode) return;
              const currentData = currentNode.data as NanoBananaNodeData;
              const histCopy = [...(currentData.imageHistory || [])];
              const entryIndex = histCopy.findIndex((h) => h.id === item.id);
              const localizedImage = saveResult.imageData || item.image;

              if (entryIndex !== -1) {
                histCopy[entryIndex] = { ...histCopy[entryIndex], image: localizedImage };
              }

              const update: Partial<NanoBananaNodeData> = { imageHistory: histCopy };
              if (currentData.outputImage === item.image || currentData.selectedHistoryIndex === entryIndex) {
                update.outputImage = localizedImage;
                update.outputImageRef = item.id;
                update.outputImageStorageStatus = "localized";
              }
              updateNodeData(node.id, update);
            })
            .catch((err) => {
              console.error("Failed to save browser-local generation:", err);
              if (isHttpMediaUrl(item.image)) {
                updateNodeData(node.id, { outputImageStorageStatus: "remote-only" });
              }
            });

          if (isHttpMediaUrl(item.image)) {
            updateNodeData(node.id, { outputImageStorageStatus: "localizing" });
          }
          trackSaveGeneration(item.id, savePromise);
        } else if (effectiveGenerationsPath) {
          const savePromise = fetch("/api/save-generation", {
            method: "POST",
            headers: { "Content-Type": "application/json" },
            body: JSON.stringify({
              directoryPath: effectiveGenerationsPath,
              image: item.image,
              prompt: finalPrompt,
              imageId: item.id,
            }),
          })
            .then((res) => res.json())
            .then(async (saveResult) => {
              if (saveResult.success && saveResult.imageId) {
                const savedImageId = String(saveResult.imageId);
                const localizedImage = await loadSavedGenerationImage(effectiveGenerationsPath, savedImageId)
                  .catch(() => null);
                const currentNode = getNodes().find((n) => n.id === node.id);
                if (!currentNode) return;

                const currentData = currentNode.data as NanoBananaNodeData;
                const histCopy = [...(currentData.imageHistory || [])];
                const entryIndex = histCopy.findIndex((h) => h.id === item.id);
                if (entryIndex !== -1) {
                  histCopy[entryIndex] = {
                    ...histCopy[entryIndex],
                    id: savedImageId,
                    image: localizedImage || histCopy[entryIndex].image,
                  };
                }

                const update: Partial<NanoBananaNodeData> = { imageHistory: histCopy };
                if (currentData.outputImage === item.image || currentData.selectedHistoryIndex === entryIndex) {
                  if (localizedImage) update.outputImage = localizedImage;
                  update.outputImageRef = savedImageId;
                  update.outputImageStorageStatus = localizedImage ? "localized" : "remote-only";
                }
                updateNodeData(node.id, update);
              }
            })
            .catch((err) => {
              console.error("Failed to save generation:", err);
              if (isHttpMediaUrl(item.image)) {
                updateNodeData(node.id, { outputImageStorageStatus: "remote-only" });
              }
            });

          if (isHttpMediaUrl(item.image)) {
            updateNodeData(node.id, { outputImageStorageStatus: "localizing" });
          }
          trackSaveGeneration(item.id, savePromise);
        }
      });
    } catch (error) {
      if (error instanceof DOMException && error.name === "AbortError") {
        throw error;
      }

      // Convert network errors to user-friendly messages
      let errorMessage = "Generation failed";
      if (error instanceof TypeError && error.message.includes("NetworkError")) {
        errorMessage = "Network error. Check your connection and try again.";
      } else if (error instanceof TypeError) {
        errorMessage = `Network error: ${error.message}`;
      } else if (error instanceof Error) {
        errorMessage = error.message;
      }

      updateNodeData(node.id, {
        status: "error",
        error: errorMessage,
      });
      throw new Error(errorMessage);
    } finally {
      await cleanupTemporaryImages(temporaryImageIds);
    }
  };

  // Synthesize a SelectedModel for the primary from legacy fields if selectedModel is missing.
  let primaryModel: SelectedModel = nodeConfig.selectedModel;
  let primaryParameters: Record<string, unknown> | undefined;
  let fallbackModel = nodeData.fallbackModel;
  let shouldFallback: ((error: unknown) => boolean) | undefined;

  if (
    primaryModel.provider === "gemini" &&
    shouldPreferNewApiWGImageModel(providerSettings, providerEnvStatus)
  ) {
    const promotedModel =
      fallbackModel && fallbackModel.provider !== "gemini"
        ? fallbackModel
        : createNewApiWGDefaultImageModel();
    primaryModel = promotedModel;
    primaryParameters = fallbackModel === promotedModel
      ? nodeData.fallbackParameters
      : nodeConfig.parameters;
    fallbackModel = undefined;
    updateNodeData(node.id, {
      selectedModel: promotedModel,
      parameters: primaryParameters ?? {},
      fallbackModel: undefined,
      fallbackParameters: undefined,
      status: nodeData.status === "error" ? "idle" : nodeData.status,
      error: nodeData.status === "error" ? null : nodeData.error,
      __usedFallback: undefined,
      __fallbackModelUsed: undefined,
      __primaryError: undefined,
      inputSchema: undefined,
    });
  }

  if (!fallbackModel && isNewApiWGDefaultImageModel(primaryModel)) {
    fallbackModel = createNewApiWGDefaultImageFallbackModel(primaryModel);
    shouldFallback = isProviderOverloadError;
  }

  await runWithFallback({
    nodeId: node.id,
    primary: primaryModel,
    primaryParameters,
    fallback: fallbackModel,
    fallbackParameters: nodeData.fallbackParameters,
    updateNodeData,
    runOnce,
    shouldFallback,
    clearOutput: { outputImage: null },
  });
}