magic-building/src/server/services/ImageGenService.ts

import { GoogleGenAI } from '@google/genai';
import mime from 'mime';
import fs from 'fs';
import path from 'path';
import { ImageGenerationOptions, ImageGenerationResult, ReferenceImage } from '../types/api';

export class ImageGenService {
  private ai: GoogleGenAI;
  private primaryModel = 'gemini-2.5-flash-image-preview';
  private fallbackModel = 'imagen-4.0-generate-001';

  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
   */
  async generateImage(options: ImageGenerationOptions): Promise<ImageGenerationResult> {
    const { prompt, filename, referenceImages } = options;
    const timestamp = new Date().toISOString();

    console.log(`[${timestamp}] Starting image generation: "${prompt.substring(0, 50)}..."`);

    try {
      // First try the primary model (Nano Banana)
      const result = await this.tryGeneration({
        model: this.primaryModel,
        config: { responseModalities: ['IMAGE', 'TEXT'] },
        prompt,
        filename,
        referenceImages,
        modelName: 'Nano Banana'
      });

      if (result.success) {
        return result;
      }

      // Fallback to Imagen 4
      console.log(`[${new Date().toISOString()}] Primary model failed, trying fallback (Imagen 4)...`);

      return await this.tryGeneration({
        model: this.fallbackModel,
        config: { responseModalities: ['IMAGE'] },
        prompt,
        filename: `${filename}_fallback`,
        referenceImages,
        modelName: 'Imagen 4'
      });

    } catch (error) {
      console.error(`[${new Date().toISOString()}] Image generation failed:`, error);
      return {
        success: false,
        model: 'none',
        error: error instanceof Error ? error.message : 'Unknown error occurred'
      };
    }
  }

  /**
   * Try generation with a specific model
   */
  private async tryGeneration(params: {
    model: string;
    config: { responseModalities: string[] };
    prompt: string;
    filename: string;
    referenceImages?: ReferenceImage[];
    modelName: string;
  }): Promise<ImageGenerationResult> {

    const { model, config, prompt, filename, referenceImages, modelName } = params;

    try {
      // Build content parts for the API request
      const contentParts: any[] = [];

      // Add reference images if provided
      if (referenceImages && referenceImages.length > 0) {
        console.log(`[${new Date().toISOString()}] Adding ${referenceImages.length} reference image(s)`);

        for (const refImage of referenceImages) {
          contentParts.push({
            inlineData: {
              mimeType: refImage.mimetype,
              data: refImage.buffer.toString('base64')
            }
          });
        }
      }

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

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

      console.log(`[${new Date().toISOString()}] Making API request to ${modelName} (${model})...`);

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

      console.log(`[${new Date().toISOString()}] Response received from ${modelName}`);

      if (response.candidates && response.candidates[0] && response.candidates[0].content) {
        const content = response.candidates[0].content;
        let generatedDescription = '';
        let savedImagePath = '';

        for (let index = 0; index < (content.parts?.length || 0); index++) {
          const part = content.parts![index];

          if (part.inlineData) {
            const fileExtension = mime.getExtension(part.inlineData.mimeType || '');
            const finalFilename = `${filename}.${fileExtension}`;
            const filepath = path.join('./src/results', finalFilename);

            console.log(`[${new Date().toISOString()}] Saving image: ${finalFilename}`);

            const buffer = Buffer.from(part.inlineData.data || '', 'base64');
            await this.saveImageFile(filepath, buffer);

            savedImagePath = filepath;

          } else if (part.text) {
            generatedDescription = part.text;
            console.log(`[${new Date().toISOString()}] Generated description: ${part.text.substring(0, 100)}...`);
          }
        }

        if (savedImagePath) {
          return {
            success: true,
            filename: path.basename(savedImagePath),
            filepath: savedImagePath,
            description: generatedDescription,
            model: modelName
          };
        }
      }

      return {
        success: false,
        model: modelName,
        error: 'No image data received from API'
      };

    } catch (error) {
      console.error(`[${new Date().toISOString()}] ${modelName} generation failed:`, error);
      return {
        success: false,
        model: modelName,
        error: error instanceof Error ? error.message : 'Generation failed'
      };
    }
  }

  /**
   * Save image buffer to file system
   */
  private async saveImageFile(filepath: string, buffer: Buffer): Promise<void> {
    return new Promise((resolve, reject) => {
      // Ensure the results directory exists
      const dir = path.dirname(filepath);
      if (!fs.existsSync(dir)) {
        fs.mkdirSync(dir, { recursive: true });
      }

      fs.writeFile(filepath, buffer, (err) => {
        if (err) {
          console.error(`[${new Date().toISOString()}] Error saving file ${filepath}:`, err);
          reject(err);
        } else {
          console.log(`[${new Date().toISOString()}] File saved successfully: ${filepath}`);
          resolve();
        }
      });
    });
  }

  /**
   * Validate reference images
   */
  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 };
  }

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