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';
import { TTILogger, TTILogEntry } from './TTILogger';

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

  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, aspectRatio, orgId, projectId, meta } = 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';
    const finalAspectRatio = aspectRatio || '1:1'; // Default to square

    // Step 1: Generate image from Gemini AI
    let generatedData: GeneratedImageData;
    let geminiParams: GeminiParams;
    try {
      const aiResult = await this.generateImageWithAI(
        prompt,
        referenceImages,
        finalAspectRatio,
        finalOrgId,
        finalProjectId,
        meta,
      );
      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[] | undefined,
    aspectRatio: string,
    orgId: string,
    projectId: string,
    meta?: { tags?: string[] },
  ): 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,
    });

    // CRITICAL: Calculate exact values before SDK call
    // These exact objects will be passed to both SDK and logger
    const contents = [
      {
        role: 'user' as const,
        parts: contentParts,
      },
    ];

    const config = {
      responseModalities: ['IMAGE', 'TEXT'],
      imageConfig: {
        aspectRatio,
      },
    };

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

    // Log TTI request BEFORE SDK call - using exact same values
    const ttiLogger = TTILogger.getInstance();
    const logEntry: TTILogEntry = {
      timestamp: new Date().toISOString(),
      orgId,
      projectId,
      prompt,
      model: this.primaryModel,
      config,
      ...(meta && { meta }),
      ...(referenceImages &&
        referenceImages.length > 0 && {
          referenceImages: referenceImages.map((img) => ({
            mimetype: img.mimetype,
            size: img.buffer.length,
            originalname: img.originalname,
          })),
        }),
    };

    ttiLogger.log(logEntry);

    try {
      // Use the EXACT same config and contents objects calculated above
      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,
    }));
  }
}