import base64
import io

from PIL import Image
import requests

def call_api_and_save_image():
    url = "http://localhost:8000/sketch-to-image"
    prompt = "A lamp with a shade on a table"
    # Request headers
    data = {
        "prompt": prompt,
        "negative_prompt": "nothing",
    }

    image_path = "run.png"
    with open(image_path, "rb") as image_file:
        files = {"image": ("run.png", image_file, "image/png")}

        response = requests.post(url, data=data, files=files)

    # Make the POST request

    # Check if the request was successful
    if response.status_code == 200:
        # Get the image URL from the response
        response_dict = response.json()

        base64_string = response_dict.get("images")[0].get("url")
        image_url = base64_string.split(",")[1]

        if image_url:
            image_data = base64.b64decode(image_url)

            # Use io.BytesIO to convert the bytes into a file-like object
            image_buffer = io.BytesIO(image_data)

            # Open the image using PIL
            image = Image.open(image_buffer)

            # Save the image to a file (e.g., "output_image.png")
            image.save("output_image.png", format="PNG")

        else:
            print("No image URL found in the response")
    else:
        print(f"API request failed. Status code: {response.status_code}")


# Call the function with the prompt
call_api_and_save_image()