#!/usr/bin/env python3
"""
Simple test for the consolidated main.py
"""

import sys
import os
import io
from PIL import Image
import base64

# Add current directory to PYTHONPATH
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)

def create_test_image():
    """Create a simple test image"""
    img = Image.new('RGB', (2048, 2048), color='lightblue')
    from PIL import ImageDraw
    draw = ImageDraw.Draw(img)
    draw.rectangle([500, 500, 1500, 1500], fill='red')
    
    buffer = io.BytesIO()
    img.save(buffer, format='PNG')
    return buffer.getvalue()

def test_preview_function():
    """Test the preview image creation"""
    # Import the function from main.py
    from main import create_preview_image_b64, get_image_info
    
    # Create test image
    test_png = create_test_image()
    print(f"Created test PNG: {len(test_png)} bytes")
    
    # Get image info
    info = get_image_info(test_png)
    print(f"Image info: {info}")
    
    # Create preview
    preview_b64 = create_preview_image_b64(test_png, target_size=512, quality=85)
    
    if preview_b64:
        print(f"✅ Preview created: {len(preview_b64)} chars")
        
        # Save preview to verify
        preview_bytes = base64.b64decode(preview_b64)
        with open('test_preview.jpg', 'wb') as f:
            f.write(preview_bytes)
        
        ratio = (len(preview_bytes) / len(test_png)) * 100
        print(f"Compression: {ratio:.1f}% (original: {len(test_png)}, preview: {len(preview_bytes)})")
        print("Preview saved as test_preview.jpg")
        return True
    else:
        print("❌ Failed to create preview")
        return False

def test_imports():
    """Test all necessary imports"""
    try:
        from main import (
            Imagen4MCPServer, 
            Imagen4ToolHandlers,
            ImageGenerationResult,
            get_tools
        )
        print("✅ All classes imported successfully")
        
        # Test tool definitions
        tools = get_tools()
        print(f"✅ {len(tools)} tools defined")
        
        return True
    except Exception as e:
        print(f"❌ Import error: {e}")
        return False

def main():
    """Main test function"""
    print("=== Testing Consolidated main.py ===\n")
    
    success = True
    
    # Test imports
    if not test_imports():
        success = False
    
    print()
    
    # Test preview function
    if not test_preview_function():
        success = False
    
    print(f"\n=== Test Results ===")
    if success:
        print("✅ All tests passed!")
        print("\nTo run the server:")
        print("1. Make sure .env is configured")
        print("2. Run: python main.py")
        print("3. Or use: run.bat")
    else:
        print("❌ Some tests failed")
    
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main())