#!/usr/bin/env python3
"""
Test the updated image generation function
"""

from ai_services import generate_image_with_huggingface

def test_image_generation():
    """Test the updated image generation function"""
    print("🧪 Testing updated image generation function...")
    
    # Test with a simple topic
    result = generate_image_with_huggingface(
        "educational diagram of ancient persia", 
        "ancient persia", 
        "worksheet"
    )
    
    if result:
        print("✅ Image generation successful!")
        print(f"Result type: {type(result)}")
        print(f"Result length: {len(result)}")
        print(f"Result starts with: {result[:50]}...")
        
        # Save the result to a file
        if result.startswith("data:image"):
            with open("test_generated_image.png", "wb") as f:
                import base64
                image_data = result.split(",")[1]  # Remove data:image/png;base64, prefix
                f.write(base64.b64decode(image_data))
            print("✅ Image saved as test_generated_image.png")
        else:
            print(f"❌ Unexpected result format: {result[:100]}")
    else:
        print("❌ Image generation failed")

if __name__ == "__main__":
    test_image_generation()