import logging

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

from bria_mcp.image_content import ImageContentMaker

logger = logging.getLogger(__name__)


async def remove_background(image: str) -> CallToolResult:
    """Remove background from an image.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.

    Returns:
        FastMCP Image object and a url for the full preview on the image
    """

    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)

    response = await client.run(endpoint="/image/edit/remove_background", payload={"image": Image(image).as_bria_api_input})
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def blur_background(image: str, scale: int = 5) -> CallToolResult:
    """Apply blur effect to an image background.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
        scale: Blur intensity (1-5, default: 5)

    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)
    scale_clamped = min(max(scale, 1), 5)
    response = await client.run(endpoint="/image/edit/blur_background", payload={"image": Image(image).as_bria_api_input, "scale": scale_clamped})
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def crop_out_foreground(image: str, padding: int | None = None) -> CallToolResult:
    """remove the background from an image and crop tightly around the foreground or remaining region of interest.
     Supports both images with and without a background.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
        padding: Optional Integer, of how much padding to add to the region of interest.

    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)
    response = await client.run(endpoint="/image/edit/crop_foreground", payload={"image": Image(image).as_bria_api_input, "padding": padding})
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def increase_resolution(image: str, scale: int = 2) -> CallToolResult:
    """Upscale image by 2x or 4x.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
        scale: Upscaling factor (2 or 4, default: 2)

    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    if scale not in [2, 4]:
        raise ValueError("Error: scale must be 2 or 4")

    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)
    response = await client.run(endpoint="/image/edit/increase_resolution", payload={"image": Image(image).as_bria_api_input, "desired_increase": scale})
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def expand_image(image: str, aspect_ratio: str | None = None, preserve_alpha: bool = False, content_moderation: bool = False) -> CallToolResult:
    """Expand an image or change its aspect ratio using generative AI.

    This function can expand images, change aspect ratios, or generate multiple aspect ratios.
    Use this tool when users want to change aspect ratio, convert to specific aspect ratios like 1:1, 16:9, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16,
    or when they mention "aspect ratio" in their request.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
        aspect_ratio: str | float. Automatically expands or crops the input image to the specified aspect ratio.
            Can be one of ["1:1", "4:3", "16:9", "3:2", "2:3", "9:16", "3:4", "4:5", "5:4"] or a float number between 0.5 and 3.0.
        preserve_alpha: (boolean, default=False). Controls whether the alpha channel values from the input image are retained in the output.
        content_moderation: (boolean, default=False) If True, apply moderation to input and output.


    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)

    payload = {
        "image": Image(image).as_bria_api_input,
        "prompt": "",
        "negative_prompt": "",
        "preserve_alpha": preserve_alpha,
        "content_moderation": content_moderation,
    }
    if aspect_ratio is not None:
        payload["aspect_ratio"] = aspect_ratio

    response = await client.run(endpoint="/image/edit/expand", payload=payload)
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def generate_background(image: str, prompt: str | None = None, reference_images: list[str] | None = None, fast: bool = False) -> CallToolResult:
    """Generate or replace an image background with a generated scene.
        Must have either prompt or reference_images!

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
        prompt: str. Optional Text description of the desired background
        reference_images: List of URL of reference images for the background
        fast: (boolean, default=True). Use fast mode for better speed/quality balance


    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)

    payload = {
        "image": Image(image).as_bria_api_input,
        "mode": "high_control",
        "fast": str(fast).lower(),
    }
    if prompt is not None:
        payload["prompt"] = prompt
    if reference_images is not None:
        payload["ref_images"] = reference_images  # type: ignore

    response = await client.run(endpoint="/image/edit/replace_background", payload=payload)
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def erase_foreground(image: str) -> CallToolResult:
    """Erase the foreground object from a provided image, while generating the area behind it.

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.

    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)

    response = await client.run(endpoint="/image/edit/erase_foreground", payload={"image": Image(image).as_bria_api_input})
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )


async def enhance_image(image: str, preserve_alpha: bool = True, content_moderation: bool = False) -> CallToolResult:
    """
    Enhance an image by generating richer details, sharper textures, and higher clarity, optionally upscaling.

    Bria /enhance defaults:
      • preserve_alpha: True
      • content_moderation: False

    Args:
        image: Image URL (preferred), file path, or base64 data. URLs provide best performance.
        preserve_alpha: (boolean, default=True). Controls whether the alpha channel values from the input image are retained in the output.
        content_moderation: (boolean, default=False) If True, apply moderation to input and output.

    Returns:
        FastMCP Image object and a url for the full preview on the image
    """
    client = BriaAsyncClient()
    # img_params = image_controller.handle_image_input(image, is_v2=True)

    response = await client.run(
        endpoint="/image/edit/enhance",
        payload={"image": Image(image).as_bria_api_input, "preserve_alpha": preserve_alpha, "content_moderation": content_moderation},
    )
    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"For full image Preview use: {response.result.image_url}"),
        ]
    )
