"""MCP public surface: three intentionally small image tools."""

from __future__ import annotations

from typing import Literal

from mcp.server import MCPServer
from mcp.server.mcpserver import Image

from .service import ImageGenerationService
from .tools.generate_guided_image import generate_guided_image as guided_adapter
from .tools.generate_image import generate_image as image_adapter
from .tools.inpaint_image import inpaint_image as inpaint_adapter

Size = Literal["square_512", "landscape_512", "portrait_512"]
GuidanceMode = Literal["depth", "openpose", "openpose_full"]


def create_server(service: ImageGenerationService) -> MCPServer:
    server = MCPServer("mcp-sd15")

    @server.tool()
    async def generate_image(
        positive_prompt: str,
        negative_prompt: str | None = None,
        size: Size | None = None,
    ) -> Image:
        """Generate one PNG from a positive prompt using the configured SD1.x model."""
        result = await image_adapter(service, positive_prompt, negative_prompt, size)
        return Image(data=result.png_data, format="png")

    @server.tool()
    async def generate_guided_image(
        positive_prompt: str,
        control_mode: GuidanceMode,
        input_image: str,
        negative_prompt: str | None = None,
        size: Size | None = None,
    ) -> Image:
        """Generate one PNG guided by Depth, OpenPose, or full OpenPose.

        input_image is a Base64 PNG/JPEG or an image data URI.
        """
        result = await guided_adapter(
            service, positive_prompt, control_mode, input_image, negative_prompt, size
        )
        return Image(data=result.png_data, format="png")

    @server.tool()
    async def inpaint_image(
        positive_prompt: str,
        input_image: str,
        mask_image: str,
        negative_prompt: str | None = None,
        size: Size | None = None,
    ) -> Image:
        """Regenerate white mask areas and return one PNG.

        Both images are Base64 PNG/JPEG values or image data URIs.
        """
        result = await inpaint_adapter(
            service, positive_prompt, input_image, mask_image, negative_prompt, size
        )
        return Image(data=result.png_data, format="png")

    return server