banatie-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 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 && { 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 && { 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) continue;

          if (part.inlineData) {
            const fileExtension = mime.getExtension(
              part.inlineData.mimeType || "",
            );
            const finalFilename = `${filename}.${fileExtension}`;
            const filepath = path.join("./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,
    }));
  }
}