from typing import Literal

from bria_client import BriaAsyncClient
from bria_client.toolkit import Image
from mcp.server.fastmcp import Context
from mcp.types import CallToolResult, TextContent

from bria_mcp.image_content import ImageContentMaker


async def generate_image(
    ctx: Context,
    image: str | None = None,
    prompt: str | None = None,
    negative_prompt: str | None = None,
    structured_prompt: dict | None = None,
    aspect_ratio: Literal["1:1", "16:9", "9:16", "4:3", "3:4"] | None = None,
    seed: int = 42,
    model_version: Literal["FIBO", "FIBO_LITE"] = "FIBO",
) -> CallToolResult:
    """
    Generate an image using Bria's FIBO model with text-to-image or image-to-image capabilities.

    FIBO is Bria's advanced image generation model that supports both text-to-image and
    image-to-image generation with precise control over the output.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
               For text-to-image generation, this can be None or an optional reference image.
        prompt: Text description of the desired output image. Guides the generation process.
        negative_prompt: Text description of what to avoid in the generated image.
        structured_prompt: Dictionary containing structured prompt parameters for more precise control.
                          Can include elements like style, composition, lighting, etc.
        aspect_ratio: Desired aspect ratio for the generated image. Common values include:
                     "1:1", "16:9", "9:16", "4:3", "3:4", etc.
        seed: Random seed for reproducible results (default 42).
        model_version: The FIBO model variant to use. Options:
                      - "FIBO": Standard FIBO model (default)
                      - "FIBO_LITE": Lightweight version for faster generation

    Returns:
        CallToolResult containing:
        - ImageContent: The generated image as base64-encoded PNG
        - TextContent: URL for full preview of the generated image
    """

    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)
    endpoint = "/image/generate"
    if model_version == "FIBO_LITE":
        endpoint = "/image/generate/lite"

    payload = {
        "model_version": "FIBO",
        "seed": seed,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "aspect_ratio": aspect_ratio,
    }
    for key, value in payload.copy().items():
        if value is None:
            payload.pop(key)
    if image is not None:
        payload["images"] = [Image(image).as_bria_api_input]
    if (structured_prompt is not None) and (structured_prompt != {}):
        payload["structured_prompt"] = structured_prompt

    status_response = await client.submit(endpoint=endpoint, payload=payload)
    status_response.raise_for_status()

    await ctx.info("Polling result (this may take a few minutes)...")
    response = await client.poll(status_response, timeout=300, interval=2)
    response.raise_for_status()

    assert response.result is not None  # Guaranteed after raise_for_status()

    return CallToolResult(
        content=[
            ImageContentMaker(response.result.image_url).to_viewable_content(),
            TextContent(type="text", text=f"Structured prompt used to build image: {response.result.structured_prompt}"),
            TextContent(type="text", text=f"For full image Preview use: {response.result.image_url}"),
        ]
    )
