import requests
import base64
from PIL import Image
import io

# API endpoint URL (adjust if running on a different host/port)
API_URL = "http://localhost:8000/predict"

# Example image URL from main.py
IMAGE_URL = "https://rail.eecs.berkeley.edu/datasets/bridge_release/raw/bridge_data_v2/datacol2_toykitchen7/drawer_pnp/01/2023-04-19_09-18-15/raw/traj_group0/traj0/images0/im_12.jpg"
TASK_TEXT = "pick up the fork"

def test_api(image_url=IMAGE_URL, task=TASK_TEXT):
    try:
        # Download image from URL
        response = requests.get(image_url, stream=True)
        response.raise_for_status()  # Check for HTTP errors
        img = Image.open(response.raw).resize((256, 256))

        # Convert image to base64
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format="JPEG")  # Save as JPEG (adjust if needed)
        img_byte_arr = img_byte_arr.getvalue()
        base64_string = base64.b64encode(img_byte_arr).decode("utf-8")

        # Prepare payload for API
        payload = {
            "image_base64": base64_string,
            "task": task
        }

        # Send POST request to API
        api_response = requests.post(API_URL, json=payload)
        api_response.raise_for_status()  # Check for API errors

        # Print the result
        result = api_response.json()
        print(f"Task: {task}")
        print(f"Image URL: {image_url}")
        print(f"Predicted Actions: {result['actions']}")

    except requests.exceptions.RequestException as e:
        print(f"Error fetching image or calling API: {e}")
    except Exception as e:
        print(f"Unexpected error: {e}")

if __name__ == "__main__":
    # Test with default values (same as main.py)
    test_api()

    # Test with a different URL and task (optional)
    # Replace with another valid URL if desired    
    print("\nTesting with another URL and task:")
    test_api(IMAGE_URL, TASK_TEXT)