#!/usr/bin/env python3
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO

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

# Test with a simple prompt
prompt = "Create a simple modern app icon with a blue background and white envelope symbol, square format, flat design, minimalist"

print(f"Testing gemini-2.0-flash-exp...")
print(f"Prompt: '{prompt}'...")

try:
    response = client.models.generate_content(
        model='gemini-2.0-flash-exp',
        contents=prompt,
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE"]
        )
    )

    # Save the Image
    if response.candidates[0].content.parts:
        for part in response.candidates[0].content.parts:
            if part.inline_data:
                image_data = part.inline_data.data
                image = Image.open(BytesIO(image_data))
                image.save("test_2_0_flash.png")
                print("✓ Success! Saved test_2_0_flash.png")
            else:
                print(f"✗ Model returned text: {part.text[:100]}")
    else:
        print("✗ No content parts in response")

except Exception as e:
    print(f"✗ Error: {type(e).__name__}: {str(e)[:300]}")