"""
Test script for Virtual Try-On API
This can be used to test the API locally or on Hugging Face Spaces
"""

import requests
import base64
from PIL import Image
import io

# API URL - change this to your Hugging Face Space URL when deployed
API_URL = "http://localhost:7860"  # Local testing
# API_URL = "https://your-username-virtual-tryon-api.hf.space"  # HF Spaces

def test_health():
    """Test health check endpoint"""
    response = requests.get(f"{API_URL}/health")
    print("Health Check:", response.json())

def image_to_base64(image_path):
    """Convert image file to base64 string"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def base64_to_image(base64_str):
    """Convert base64 string to PIL Image"""
    img_data = base64.b64decode(base64_str)
    return Image.open(io.BytesIO(img_data))

def test_tryon_with_files(person_image_path, clothing_image_path):
    """Test virtual try-on with file upload"""
    print("\n🧪 Testing Virtual Try-On (File Upload)...")
    
    with open(person_image_path, 'rb') as person_file, \
         open(clothing_image_path, 'rb') as clothing_file:
        
        files = {
            'person_image': person_file,
            'clothing_image': clothing_file,
        }
        
        data = {
            'return_format': 'base64',
            'num_steps': 30,  # Faster for testing
        }
        
        response = requests.post(f"{API_URL}/tryon", files=files, data=data)
        
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Success! Processing time: {result['processing_time']:.2f}s")
            
            # Save the result
            generated_image = base64_to_image(result['image'])
            generated_image.save("test_result.png")
            print("📁 Result saved as 'test_result.png'")
            
            return result
        else:
            print(f"❌ Error: {response.status_code}")
            print(response.text)

def test_tryon_with_base64(person_image_path, clothing_image_path):
    """Test virtual try-on with base64 encoding"""
    print("\n🧪 Testing Virtual Try-On (Base64)...")
    
    person_base64 = image_to_base64(person_image_path)
    clothing_base64 = image_to_base64(clothing_image_path)
    
    data = {
        'person_image_base64': person_base64,
        'clothing_image_base64': clothing_base64,
        'num_steps': 30,  # Faster for testing
    }
    
    response = requests.post(f"{API_URL}/tryon-base64", data=data)
    
    if response.status_code == 200:
        result = response.json()
        print(f"✅ Success! Processing time: {result['processing_time']:.2f}s")
        
        # Save the result
        generated_image = base64_to_image(result['image'])
        generated_image.save("test_result_base64.png")
        print("📁 Result saved as 'test_result_base64.png'")
        
        return result
    else:
        print(f"❌ Error: {response.status_code}")
        print(response.text)

if __name__ == "__main__":
    print("🚀 Starting API Tests\n")
    
    # Test health check
    test_health()
    
    # Test with your own images
    # Replace these paths with your test images
    person_image = "test_person.jpg"
    clothing_image = "test_clothing.jpg"
    
    # Uncomment to test (make sure you have test images)
    # test_tryon_with_files(person_image, clothing_image)
    # test_tryon_with_base64(person_image, clothing_image)
    
    print("\n✅ Tests completed!")