from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import os

# Initialize the Client
client = genai.Client(api_key="AIzaSyB-nR0nkeftKrd42NrNIDcFCj3yFP8JLtw")

# Define the prompt
prompt = "A simple modern app icon with a blue background and white envelope symbol, square format, flat design"

print(f"Generating image for: '{prompt}'...")

# Call the API
try:
    response = client.models.generate_images(
        model='gemini-2.5-flash-image',
        prompt=prompt,
        config=types.GenerateImagesConfig(
            number_of_images=1,
            aspect_ratio="1:1",
            safety_filter_level="BLOCK_LOW_AND_ABOVE",
            person_generation="ALLOW_ADULT"
        )
    )

    # Handle the response
    for i, generated_image in enumerate(response.generated_images):
        # Convert raw bytes to an image
        image = Image.open(BytesIO(generated_image.image.image_bytes))

        # Save to disk
        filename = f"test_generated_image_{i}.png"
        image.save(filename)
        print(f"Success! Saved {filename}")

except Exception as e:
    print(f"Error occurred: {type(e).__name__}: {e}")
    import traceback
    traceback.print_exc()