// To run this code you need to install the following dependencies:
// npm install @google/genai mime dotenv
// npm install -D @types/node

import 'dotenv/config';
import {
  GoogleGenAI,
} from '@google/genai';
import mime from 'mime';
import { writeFile } from 'fs';

function saveBinaryFile(fileName: string, content: Buffer) {
  writeFile(fileName, content, (err) => {
    if (err) {
      console.error(`Error writing file ${fileName}:`, err);
      return;
    }
    console.log(`File ${fileName} saved to file system.`);
  });
}

async function main() {
  console.log('Loading environment variables...');
  console.log('API Key loaded:', process.env.GEMINI_API_KEY ? 'Yes' : 'No');

  if (!process.env.GEMINI_API_KEY) {
    console.error('GEMINI_API_KEY environment variable is not set!');
    process.exit(1);
  }

  console.log('Creating GoogleGenAI instance...');
  const ai = new GoogleGenAI({
    apiKey: process.env.GEMINI_API_KEY,
  });

  // Note: Using gemini-2.5-flash-image-preview (Nano Banana) as primary model

  const config = {
    responseModalities: ['IMAGE', 'TEXT'],
  };
  const model = 'gemini-2.5-flash-image-preview';
  const contents = [
    {
      role: 'user',
      parts: [
        {
          text: `Generate a fantasy warrior with magical sword`,
        },
      ],
    },
  ];

  console.log('Making API request for image generation...');
  try {
    const response = await ai.models.generateContent({
      model,
      config,
      contents,
    });

    console.log('Response received successfully');

    if (response.candidates && response.candidates[0] && response.candidates[0].content) {
      const content = response.candidates[0].content;

      content.parts?.forEach((part, index) => {
        if (part.inlineData) {
          const fileName = `fantasy_warrior_primary_${index}`;
          const inlineData = part.inlineData;
          const fileExtension = mime.getExtension(inlineData.mimeType || '');
          console.log(`Saving image: ${fileName}.${fileExtension}`);
          const buffer = Buffer.from(inlineData.data || '', 'base64');
          saveBinaryFile(`${fileName}.${fileExtension}`, buffer);
        } else if (part.text) {
          console.log('Description:', part.text);
        }
      });
    }
    console.log('Image generation complete!');
  } catch (error) {
    console.error('Primary model failed:', error instanceof Error ? error.message : error);

    // Try fallback model (Imagen 4)
    console.log('Trying fallback model (Imagen 4)...');
    try {
      const altResponse = await ai.models.generateContent({
        model: 'imagen-4.0-generate-001',
        config: {
          responseModalities: ['IMAGE']
        },
        contents,
      });

      console.log('Fallback model succeeded!');

      if (altResponse.candidates && altResponse.candidates[0] && altResponse.candidates[0].content) {
        const content = altResponse.candidates[0].content;

        content.parts?.forEach((part, index) => {
          if (part.inlineData) {
            const fileName = `fantasy_warrior_fallback_${index}`;
            const inlineData = part.inlineData;
            const fileExtension = mime.getExtension(inlineData.mimeType || '');
            console.log(`Saving fallback image: ${fileName}.${fileExtension}`);
            const buffer = Buffer.from(inlineData.data || '', 'base64');
            saveBinaryFile(`${fileName}.${fileExtension}`, buffer);
          } else if (part.text) {
            console.log('Fallback description:', part.text);
          }
        });
      }
    } catch (altError) {
      console.error('Both models failed. Please try again later.');
      console.error('Fallback error:', altError instanceof Error ? altError.message : altError);
    }
  }
}

main();
