from mcp.server.fastmcp import FastMCP
import json
import sys
import io
import time
from gradio_client import Client

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")

mcp = FastMCP("huggingface_spaces_image_display")


@mcp.tool()
async def generate_image(prompt: str, width: int = 512, height: int = 512) -> str:
    """Generate an image using SanaSprint model.

    Args:
        prompt: Text prompt describing the image to generate
        width: Image width (default: 512)
        height: Image height (default: 512)
    """
    client = Client("https://ysharma-sanasprint.hf.space/")

    try:
        result = client.predict(
            prompt, "0.6B", 0, True, width, height, 4.0, 2, api_name="/infer"
        )

        if isinstance(result, list) and len(result) >= 1:
            image_data = result[0]
            if isinstance(image_data, dict) and "url" in image_data:
                return json.dumps(
                    {
                        "type": "image",
                        "url": image_data["url"],
                        "message": f"Generated image for prompt: {prompt}",
                    }
                )

        return json.dumps({"type": "error", "message": "Failed to generate image"})

    except Exception as e:
        return json.dumps(
            {"type": "error", "message": f"Error generating image: {str(e)}"}
        )


if __name__ == "__main__":
    mcp.run(transport="stdio")