"""
Minimal SD Agent Example

Demonstrates how to integrate Stable Diffusion image generation
into a GAIA agent using the Lemonade Server endpoint.
"""

from pathlib import Path
import base64
import requests
from gaia.agents.base import Agent
from gaia.agents.base.tools import tool


class SimpleSDAgent(Agent):
    """A minimal agent that can generate images from text descriptions."""

    def __init__(
        self,
        base_url: str = "http://localhost:8000",
        output_dir: str = "./generated_images",
        **kwargs
    ):
        super().__init__(**kwargs)
        self.sd_endpoint = f"{base_url}/api/v1/images/generations"
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._register_tools()

    def _get_system_prompt(self) -> str:
        return """You are an image generation assistant.

When the user asks you to create or generate an image, use the generate_image tool.
Enhance simple prompts with artistic details for better results.

Example enhancements:
- "a cat" → "a fluffy orange cat, soft lighting, detailed fur, photorealistic"
- "sunset" → "vibrant sunset over ocean, golden hour, dramatic clouds, 4k"
"""

    def _register_tools(self):
        @tool
        def generate_image(
            prompt: str,
            model: str = "SD-Turbo",
            size: str = "512x512",
            steps: int = 4
        ) -> dict:
            """
            Generate an image from a text prompt using Stable Diffusion.

            Args:
                prompt: Text description of the image to generate
                model: SD model to use (SD-Turbo or SDXL-Turbo)
                size: Image dimensions (512x512, 768x768, 1024x1024)
                steps: Number of inference steps (4 for Turbo models)

            Returns:
                Dict with image_path and generation details
            """
            # Call SD endpoint
            response = requests.post(
                self.sd_endpoint,
                json={
                    "prompt": prompt,
                    "model": model,
                    "size": size,
                    "n": 1,
                    "response_format": "b64_json"
                },
                timeout=60
            )
            response.raise_for_status()

            # Decode and save image
            data = response.json()
            image_b64 = data["data"][0]["b64_json"]
            image_bytes = base64.b64decode(image_b64)

            # Generate filename from prompt
            safe_name = "".join(c if c.isalnum() else "_" for c in prompt[:30])
            image_path = self.output_dir / f"{safe_name}.png"
            image_path.write_bytes(image_bytes)

            return {
                "image_path": str(image_path),
                "prompt": prompt,
                "model": model,
                "size": size
            }

        self.register_tool(generate_image)


# CLI usage
if __name__ == "__main__":
    agent = SimpleSDAgent()
    agent.run("Generate an image of a dragon perched on a mountain cliff at sunset")

"""Test the minimal SD agent example."""

import pytest
from pathlib import Path


@pytest.fixture
def sd_agent():
    from examples.sd_agent_minimal import SimpleSDAgent
    return SimpleSDAgent(output_dir="./test_output")


@pytest.mark.integration
@pytest.mark.requires_lemonade
def test_generate_image(sd_agent, tmp_path):
    """Test basic image generation."""
    sd_agent.output_dir = tmp_path

    result = sd_agent.tools["generate_image"](
        prompt="a simple red circle",
        model="SD-Turbo",
        size="512x512"
    )

    assert Path(result["image_path"]).exists()
    assert result["model"] == "SD-Turbo"


def test_agent_has_tool(sd_agent):
    """Test that generate_image tool is registered."""
    assert "generate_image" in sd_agent.tools