#!/usr/bin/env python3
"""
Test image generation with just a few images first
"""

import os
import sys

# Set the API key from the .env file
os.environ['OPENAI_API_KEY'] = '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'

# Add script directory to path
sys.path.insert(0, '/var/www/twin-digital-media/public_html/_sites/cleankitchens/production/scripts')

from generate_stock_images import StockImageGenerator

def test_generation():
    """Generate just a few test images"""
    generator = StockImageGenerator()
    
    # Create directories
    generator.create_directories()
    
    print("\n=== Testing with 3 images ===")
    
    # Test prompts - one from each category
    test_prompts = [
        # Rodent violation
        ("Dark industrial kitchen corner with visible mouse droppings near food prep area, yellow caution tape in foreground, Professional photography, industrial commercial kitchen setting, dramatic lighting, high contrast, cinematic quality, photorealistic, no people, no text, no logos", "rodent", 1),
        
        # Closure
        ("Restaurant entrance with official red closure notice on door, dramatic evening lighting, Professional photography, industrial commercial kitchen setting, dramatic lighting, high contrast, cinematic quality, photorealistic, no people, no text, no logos", "closure", 1),
        
        # Daily summary
        ("Clipboard with inspection checklist on stainless steel surface, dramatic overhead lighting, Professional photography, industrial commercial kitchen setting, dramatic lighting, high contrast, cinematic quality, photorealistic, no people, no text, no logos", "daily_summary", 1)
    ]
    
    results = []
    for prompt, category, index in test_prompts:
        print(f"\nGenerating {category}_{index}...")
        result = generator.generate_image(prompt, category, index)
        if result:
            results.append(result)
            print(f"✓ Generated: {result['filepath']}")
    
    print(f"\n=== Test Complete ===")
    print(f"Generated: {generator.stats['generated']} images")
    print(f"Failed: {generator.stats['failed']} images")
    print(f"Cost: ${generator.stats['total_cost']:.2f}")
    
    return results

if __name__ == "__main__":
    print("This will generate 3 test images (cost ~$0.24)")
    response = input("Proceed? (y/n): ")
    
    if response.lower() == 'y':
        test_generation()
    else:
        print("Cancelled")