import json
from typing import Any, Literal

from crewai.tools import BaseTool, EnvVar
from openai import OpenAI
from pydantic import BaseModel, Field


class ImagePromptSchema(BaseModel):
    """Input for Dall-E Tool."""

    image_description: str = Field(
        description="Description of the image to be generated by Dall-E."
    )


class DallETool(BaseTool):
    name: str = "Dall-E Tool"
    description: str = "Generates images using OpenAI's Dall-E model."
    args_schema: type[BaseModel] = ImagePromptSchema

    model: str = "dall-e-3"
    size: (
        Literal[
            "auto",
            "1024x1024",
            "1536x1024",
            "1024x1536",
            "256x256",
            "512x512",
            "1792x1024",
            "1024x1792",
        ]
        | None
    ) = "1024x1024"
    quality: Literal["standard", "hd", "low", "medium", "high", "auto"] | None = (
        "standard"
    )
    n: int = 1

    env_vars: list[EnvVar] = Field(
        default_factory=lambda: [
            EnvVar(
                name="OPENAI_API_KEY",
                description="API key for OpenAI services",
                required=True,
            ),
        ]
    )

    def _run(self, **kwargs: Any) -> str:
        client = OpenAI()

        image_description = kwargs.get("image_description")

        if not image_description:
            return "Image description is required."

        response = client.images.generate(
            model=self.model,
            prompt=image_description,
            size=self.size,
            quality=self.quality,
            n=self.n,
        )

        if not response or not response.data:
            return "Failed to generate image."

        return json.dumps(
            {
                "image_url": response.data[0].url,
                "image_description": response.data[0].revised_prompt,
            }
        )