#!/usr/bin/env python3
"""
Test script for ColorCraft SDXL Space
Verifies that your deployed Space is working correctly
"""

import requests
import base64
import json
from PIL import Image
import io
import time

def create_test_image():
    """Create a simple test image"""
    from PIL import Image, ImageDraw
    
    # Create a simple test image with shapes
    img = Image.new('RGB', (512, 512), 'white')
    draw = ImageDraw.Draw(img)
    
    # Draw some shapes for testing
    draw.rectangle([100, 100, 400, 400], outline='black', width=3)
    draw.ellipse([150, 150, 350, 350], outline='blue', width=2)
    draw.line([200, 200, 300, 300], fill='red', width=4)
    
    return img

def test_space_api(space_url, test_image_path=None):
    """Test the deployed ColorCraft SDXL Space"""
    
    print(f"🧪 Testing ColorCraft SDXL Space: {space_url}")
    
    # Use provided image or create test image
    if test_image_path:
        try:
            image = Image.open(test_image_path)
            print(f"📸 Using test image: {test_image_path}")
        except Exception as e:
            print(f"❌ Could not load image {test_image_path}: {e}")
            return False
    else:
        image = create_test_image()
        print("📸 Using generated test image")
    
    # Convert image to base64
    img_buffer = io.BytesIO()
    image.save(img_buffer, format='PNG')
    img_buffer.seek(0)
    image_b64 = base64.b64encode(img_buffer.getvalue()).decode()
    
    # Prepare API request
    payload = {
        "data": [
            f"data:image/png;base64,{image_b64}",  # image
            "clean",                              # style
            0.7,                                  # detail_level
            True                                  # noise_reduction
        ]
    }
    
    try:
        print("🚀 Sending request to Space...")
        start_time = time.time()
        
        response = requests.post(
            f"{space_url}/api/predict",
            json=payload,
            timeout=120,  # 2 minute timeout for SDXL
            headers={"Content-Type": "application/json"}
        )
        
        processing_time = time.time() - start_time
        print(f"⏱️ Request completed in {processing_time:.2f} seconds")
        
        if response.status_code == 200:
            result = response.json()
            print("✅ Space responded successfully!")
            
            # Check if we got a valid result
            if "data" in result and len(result["data"]) >= 2:
                result_image_data = result["data"][0]
                status_message = result["data"][1]
                
                print(f"📋 Status: {status_message}")
                
                if result_image_data:
                    # Try to decode and save result
                    try:
                        if isinstance(result_image_data, str):
                            if result_image_data.startswith('data:image'):
                                image_data = result_image_data.split(',')[1]
                            else:
                                image_data = result_image_data
                            
                            # Decode and save
                            image_bytes = base64.b64decode(image_data)
                            result_image = Image.open(io.BytesIO(image_bytes))
                            
                            # Save result
                            output_path = "test_result.png"
                            result_image.save(output_path)
                            print(f"🎨 Result saved as: {output_path}")
                            print(f"📏 Result size: {result_image.size}")
                            
                            return True
                        
                        elif isinstance(result_image_data, dict) and "url" in result_image_data:
                            print(f"🔗 Result URL: {result_image_data['url']}")
                            return True
                            
                    except Exception as e:
                        print(f"⚠️ Could not process result image: {e}")
                        print(f"Raw result type: {type(result_image_data)}")
                        return False
                else:
                    print("⚠️ No image data in response")
                    return False
            else:
                print("⚠️ Unexpected response format")
                print(f"Response: {json.dumps(result, indent=2)}")
                return False
        
        else:
            print(f"❌ Space returned status {response.status_code}")
            print(f"Response: {response.text}")
            return False
            
    except requests.exceptions.Timeout:
        print("⏰ Request timed out - Space might be loading or overloaded")
        return False
    except Exception as e:
        print(f"❌ Error testing Space: {e}")
        return False

def main():
    print("🧪 ColorCraft SDXL Space Tester")
    print("=" * 40)
    
    # Get Space URL
    space_url = input("Enter your Space URL (e.g., https://username-colorcraft-sdxl.hf.space): ").strip()
    
    if not space_url:
        print("❌ Space URL is required!")
        return
    
    # Remove trailing slash
    space_url = space_url.rstrip('/')
    
    # Optional test image
    test_image = input("Enter path to test image (or press Enter for generated test): ").strip()
    test_image = test_image if test_image else None
    
    # Run test
    success = test_space_api(space_url, test_image)
    
    if success:
        print("\n🎉 Space Test PASSED!")
        print("✅ Your SDXL Space is working correctly")
        print("🚀 Ready for production use!")
    else:
        print("\n❌ Space Test FAILED!")
        print("🔧 Check your Space logs and configuration")
        print("💡 Try waiting a few minutes if the Space is still building")

if __name__ == "__main__":
    main()