import io
from PIL import Image

@tool
def generate_image(prompt: str) -> AgentImage:
    """
    Generates an image from a text prompt using Stable Diffusion XL (SDXL).
    Args:
        prompt: Text prompt describing the image to generate.
    Returns:
        AgentImage object containing the generated image.
    """
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        raise ValueError("HF_TOKEN not set in Space secrets")

    api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
    headers = {"Authorization": f"Bearer {hf_token}"}
    
    # It's safer to use the 'data' or 'json' param correctly
    response = requests.post(api_url, headers=headers, json={"inputs": prompt})
    
    if response.status_code != 200:
        raise Exception(f"API Error: {response.status_code} - {response.text}")

    # Convert bytes to an actual Image object
    image_bytes = response.content
    image = Image.open(io.BytesIO(image_bytes))
    
    return image # smolagents will automatically wrap a PIL image into AgentImage