"""
Example script for using the Transparent Background Remover Inference Endpoint.

This script:
  - Loads three sample images (sample1.png, sample2.png, sample3.png)
  - Encodes them as base64 data URIs
  - Builds a JSON payload and sends a POST request to the endpoint
  - Receives the processed images (base64-encoded) and saves them as PNG files

Requirements:
  - Python 3.10+
  - requests

Usage:
  python example.py
"""

import requests
import base64


def encode_image_to_base64(filename):
    """
    Encode an image file to a base64 data URI.

    Args:
        filename (str): Path to the image file.

    Returns:
        str: Base64-encoded string with data URI prefix.
    """
    with open(filename, "rb") as f:
        image_bytes = f.read()
    encoded = base64.b64encode(image_bytes).decode("utf-8")
    # Assuming PNG format; adjust MIME type if using another format.
    data_uri = f"data:image/png;base64,{encoded}"
    return data_uri


def save_base64_to_file(data_uri, output_filename):
    """
    Decode a base64 data URI and save it as an image file.

    Args:
        data_uri (str): The data URI containing the base64-encoded image.
        output_filename (str): The output filename to save the image.
    """
    header, encoded = data_uri.split(",", 1)
    image_data = base64.b64decode(encoded)
    with open(output_filename, "wb") as f:
        f.write(image_data)
    print(f"Saved {output_filename}")


def main():
    # Replace with your actual endpoint URL and API token (if required)
    endpoint_url = "https://endpoints.huggingface.co/your-endpoint-url"
    headers = {
        "Authorization": "Bearer YOUR_API_TOKEN",
        "Content-Type": "application/json"
    }

    # Encode three sample images.
    sample_images = [
        encode_image_to_base64("bear1.png"),
        encode_image_to_base64("bear2.png"),
        encode_image_to_base64("bear3.png")
    ]

    # Build the payload according to the handler's expected format.
    payload = {
        "inputs": {
            "images": sample_images,
            "output_type": "rgba",  # Choose "rgba", "map", "green", "blur", or "overlay"
            "threshold": 0.1,
            "reverse": False        # Set to True if you want to remove the foreground instead
        }
    }

    print("Sending request to endpoint...")
    response = requests.post(endpoint_url, json=payload, headers=headers)

    if response.status_code != 200:
        print("Error:", response.status_code, response.text)
        return

    result = response.json()
    print("Processing Times:\n", result.get("processing_times", "No timing info provided."))

    # Save each processed image to a file.
    images = result.get("images", [])
    for idx, data_uri in enumerate(images):
        save_base64_to_file(data_uri, f"output_{idx+1}.png")


if __name__ == "__main__":
    main()