banatie-service/apps/api-service/src/services/ImageGenService.ts

import { GoogleGenAI } from "@google/genai";
// eslint-disable-next-line @typescript-eslint/no-var-requires
const mime = require("mime") as any;
import {
  ImageGenerationOptions,
  ImageGenerationResult,
  ReferenceImage,
  GeneratedImageData,
  GeminiParams,
} from "../types/api";
import { StorageFactory } from "./StorageFactory";

export class ImageGenService {
  private ai: GoogleGenAI;
  private primaryModel = "gemini-2.5-flash-image-preview";

  constructor(apiKey: string) {
    if (!apiKey) {
      throw new Error("Gemini API key is required");
    }
    this.ai = new GoogleGenAI({ apiKey });
  }

  /**
   * Generate an image from text prompt with optional reference images
   * This method separates image generation from storage for clear error handling
   */
  async generateImage(
    options: ImageGenerationOptions,
  ): Promise<ImageGenerationResult> {
    const { prompt, filename, referenceImages, orgId, projectId } = options;

    // Use default values if not provided
    const finalOrgId = orgId || process.env["DEFAULT_ORG_ID"] || "default";
    const finalProjectId =
      projectId || process.env["DEFAULT_PROJECT_ID"] || "main";

    // Step 1: Generate image from Gemini AI
    let generatedData: GeneratedImageData;
    let geminiParams: GeminiParams;
    try {
      const aiResult = await this.generateImageWithAI(prompt, referenceImages);
      generatedData = aiResult.generatedData;
      geminiParams = aiResult.geminiParams;
    } catch (error) {
      // Generation failed - return explicit error
      return {
        success: false,
        model: this.primaryModel,
        error:
          error instanceof Error ? error.message : "Image generation failed",
        errorType: "generation",
      };
    }

    // Step 2: Save generated image to storage
    try {
      const finalFilename = `${filename}.${generatedData.fileExtension}`;
      const storageService = await StorageFactory.getInstance();
      const uploadResult = await storageService.uploadFile(
        finalOrgId,
        finalProjectId,
        "generated",
        finalFilename,
        generatedData.buffer,
        generatedData.mimeType,
      );

      if (uploadResult.success) {
        return {
          success: true,
          filename: uploadResult.filename,
          filepath: uploadResult.path,
          url: uploadResult.url,
          model: this.primaryModel,
          geminiParams,
          ...(generatedData.description && {
            description: generatedData.description,
          }),
        };
      } else {
        // Storage failed but image was generated
        return {
          success: false,
          model: this.primaryModel,
          geminiParams,
          error: `Image generated successfully but storage failed: ${uploadResult.error || "Unknown storage error"}`,
          errorType: "storage",
          generatedImageData: generatedData,
          ...(generatedData.description && {
            description: generatedData.description,
          }),
        };
      }
    } catch (error) {
      // Storage exception - image was generated but couldn't be saved
      return {
        success: false,
        model: this.primaryModel,
        geminiParams,
        error: `Image generated successfully but storage failed: ${error instanceof Error ? error.message : "Unknown storage error"}`,
        errorType: "storage",
        generatedImageData: generatedData,
        ...(generatedData.description && {
          description: generatedData.description,
        }),
      };
    }
  }

  /**
   * Generate image using Gemini AI - isolated from storage logic
   * @throws Error if generation fails
   */
  private async generateImageWithAI(
    prompt: string,
    referenceImages?: ReferenceImage[],
  ): Promise<{ generatedData: GeneratedImageData; geminiParams: GeminiParams }> {
    const contentParts: any[] = [];

    // Add reference images if provided
    if (referenceImages && referenceImages.length > 0) {
      for (const refImage of referenceImages) {
        contentParts.push({
          inlineData: {
            mimeType: refImage.mimetype,
            data: refImage.buffer.toString("base64"),
          },
        });
      }
    }

    // Add text prompt
    contentParts.push({
      text: prompt,
    });

    const contents = [
      {
        role: "user" as const,
        parts: contentParts,
      },
    ];

    const config = { responseModalities: ["IMAGE", "TEXT"] };

    // Capture Gemini SDK parameters for debugging
    const geminiParams: GeminiParams = {
      model: this.primaryModel,
      config,
      contentsStructure: {
        role: "user",
        partsCount: contentParts.length,
        hasReferenceImages: !!(referenceImages && referenceImages.length > 0),
      },
    };

    try {
      const response = await this.ai.models.generateContent({
        model: this.primaryModel,
        config,
        contents,
      });

      // Parse response
      if (
        !response.candidates ||
        !response.candidates[0] ||
        !response.candidates[0].content
      ) {
        throw new Error("No response received from Gemini AI");
      }

      const content = response.candidates[0].content;
      let generatedDescription: string | undefined;
      let imageData: { buffer: Buffer; mimeType: string } | null = null;

      // Extract image data and description from response
      for (const part of content.parts || []) {
        if (part.inlineData) {
          const buffer = Buffer.from(part.inlineData.data || "", "base64");
          const mimeType = part.inlineData.mimeType || "image/png";
          imageData = { buffer, mimeType };
        } else if (part.text) {
          generatedDescription = part.text;
        }
      }

      if (!imageData) {
        throw new Error("No image data received from Gemini AI");
      }

      const fileExtension = mime.getExtension(imageData.mimeType) || "png";

      const generatedData: GeneratedImageData = {
        buffer: imageData.buffer,
        mimeType: imageData.mimeType,
        fileExtension,
        ...(generatedDescription && { description: generatedDescription }),
      };

      return {
        generatedData,
        geminiParams,
      };
    } catch (error) {
      // Re-throw with clear error message
      if (error instanceof Error) {
        throw new Error(`Gemini AI generation failed: ${error.message}`);
      }
      throw new Error("Gemini AI generation failed: Unknown error");
    }
  }

  static validateReferenceImages(files: Express.Multer.File[]): {
    valid: boolean;
    error?: string;
  } {
    if (files.length > 3) {
      return { valid: false, error: "Maximum 3 reference images allowed" };
    }

    const allowedTypes = ["image/png", "image/jpeg", "image/jpg", "image/webp"];
    const maxSize = 5 * 1024 * 1024; // 5MB

    for (const file of files) {
      if (!allowedTypes.includes(file.mimetype)) {
        return {
          valid: false,
          error: `Unsupported file type: ${file.mimetype}. Allowed: PNG, JPEG, WebP`,
        };
      }

      if (file.size > maxSize) {
        return {
          valid: false,
          error: `File ${file.originalname} is too large. Maximum size: 5MB`,
        };
      }
    }

    return { valid: true };
  }

  static convertFilesToReferenceImages(
    files: Express.Multer.File[],
  ): ReferenceImage[] {
    return files.map((file) => ({
      buffer: file.buffer,
      mimetype: file.mimetype,
      originalname: file.originalname,
    }));
  }
}