import requests
from PIL import Image
import io
import sys
import json

def test_local_server(image_path=None):
    """
    Test the local Gradio server by sending a request to generate embeddings for an image
    
    Args:
        image_path: Path to the image file. If None, a test URL will be used.
    """
    # Local server URL (default Gradio port)
    server_url = "http://localhost:7860/api/predict"
    
    if image_path:
        try:
            # Load the image
            image = Image.open(image_path)
            
            # Convert image to bytes
            img_byte_arr = io.BytesIO()
            image.save(img_byte_arr, format=image.format if image.format else 'PNG')
            img_byte_arr = img_byte_arr.getvalue()
            
            # Prepare the request
            files = {
                'data': ('image.png', img_byte_arr, 'image/png')
            }
            
            # Send the request
            response = requests.post(server_url, files=files)
            
        except Exception as e:
            print(f"Error loading image: {str(e)}")
            return
    else:
        # Use a test image URL
        test_image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/bert-architects.png"
        print(f"Using test image URL: {test_image_url}")
        
        # Download the image
        response = requests.get(test_image_url)
        image = Image.open(io.BytesIO(response.content))
        
        # Convert image to bytes
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()
        
        # Prepare the request
        files = {
            'data': ('image.png', img_byte_arr, 'image/png')
        }
        
        # Send the request
        response = requests.post(server_url, files=files)
    
    print("Sending request to local Gradio server...")
    
    try:
        # Check if the request was successful
        if response.status_code == 200:
            # Parse the response
            result = response.json()
            
            if "error" in result:
                print(f"Error from server: {result['error']}")
            else:
                # Extract the embedding from the response
                embedding_data = result['data'][0]
                embedding_dim = result['data'][1]
                
                print("✅ Test successful!")
                print(f"Embedding dimension: {embedding_dim}")
                if embedding_data:
                    print(f"First 10 values of embedding: {embedding_data['embedding'][:10]}...")
                else:
                    print("No embedding data returned")
        else:
            print(f"❌ Error: HTTP {response.status_code}")
            print(response.text)
    
    except Exception as e:
        print(f"❌ Error connecting to server: {str(e)}")
        print("Make sure the server is running with 'python app.py'")

if __name__ == "__main__":
    # Check if an image path was provided
    if len(sys.argv) > 1:
        image_path = sys.argv[1]
        print(f"Testing with image: {image_path}")
        test_local_server(image_path)
    else:
        print("No image path provided, using test URL")
        test_local_server()