# tools/text_to_image.py
import os
from huggingface_hub import InferenceClient
from smolagents import Tool


class TextToImageTool(Tool):
    name = "image_generator"
    description = "Generate an image from a text prompt."
    inputs = {
        "prompt": {
            "type": "string",
            "description": "Text prompt for the image. Add style hints like photorealistic, high-res, etc.",
        }
    }
    output_type = "image"

    def __init__(self, model_id: str = "stabilityai/stable-diffusion-xl-base-1.0"):
        super().__init__()
        token = (
            os.getenv("HF_TOKEN")
            or os.getenv("HUGGINGFACEHUB_API_TOKEN")
            or os.getenv("HF_API_TOKEN")
        )
        self.client = InferenceClient(model=model_id, token=token)

    def forward(self, prompt: str):
        return self.client.text_to_image(prompt)