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

def test_omniparser_api(api_url, image_path):
    """
    Test the OmniParser API by sending an image and displaying the results
    
    Args:
        api_url: URL of the OmniParser API
        image_path: Path to the image file to test
    """
    print(f"Testing OmniParser API at: {api_url}")
    print(f"Using image: {image_path}")
    
    # Upload the image file
    files = {'image': open(image_path, 'rb')}
    
    try:
        # Send request to the API
        response = requests.post(api_url, files=files)
        
        # Check if the request was successful
        if response.status_code == 200:
            # Parse the JSON response
            result = response.json()
            
            if "error" in result:
                print(f"Error: {result['error']}")
                return
            
            # Display the parsed elements
            print("\nParsed Elements:")
            print("---------------")
            
            elements = result["elements"]
            for element in elements:
                print(f"Element {element['id']}:")
                print(f"  Text: {element['text']}")
                print(f"  Caption: {element['caption']}")
                print(f"  Coordinates: {element['coordinates']}")
                print(f"  Interactable: {element['is_interactable']}")
                print(f"  Confidence: {element['confidence']}")
                print("---")
            
            # Display the visualization
            if "visualization" in result:
                print("\nVisualization received. Saving to 'visualization.png'")
                img_data = base64.b64decode(result["visualization"])
                img = Image.open(io.BytesIO(img_data))
                img.save("visualization.png")
                print("Visualization saved to 'visualization.png'")
            
        else:
            print(f"Error: Request failed with status code {response.status_code}")
            print(response.text)
    
    except Exception as e:
        print(f"Error: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test the OmniParser API")
    parser.add_argument("--api_url", type=str, required=True, help="URL of the OmniParser API")
    parser.add_argument("--image_path", type=str, required=True, help="Path to the image file to test")
    
    args = parser.parse_args()
    test_omniparser_api(args.api_url, args.image_path)