"""Tools for querying Midjourney images."""

import io
from typing import Any

from mcp.server.fastmcp.utilities.types import Image
from PIL import Image as PILImage

from ..client import ImagineAPIClient
from ..models import ImageStatus


async def midjourney_get_status(image_id: str) -> dict[str, Any]:
    """Get image generation status.

    Args:
        image_id: Image ID to query

    Returns:
        Dictionary containing image status info:
        - id: Image ID
        - prompt: The prompt used
        - status: Current status (pending, in-progress, completed, failed)
        - progress: Progress percentage (0-100, only when in-progress)
        - url: Completed image URL (only when completed)
        - upscaled_urls: List of upscaled image URLs (only when completed)
        - error: Error message (only when failed)

    Raises:
        ImagineAPIError: If API call fails or image not found
    """
    async with ImagineAPIClient() as client:
        image = await client.get_image(image_id)

    result: dict[str, Any] = {
        "id": image.id,
        "prompt": image.prompt,
        "status": image.status.value,
    }

    if image.status == ImageStatus.IN_PROGRESS and image.progress is not None:
        result["progress"] = image.progress

    if image.status == ImageStatus.COMPLETED:
        result["url"] = image.result or image.url
        if image.upscaled_urls:
            result["upscaled_urls"] = image.upscaled_urls

    if image.status == ImageStatus.FAILED and image.error:
        result["error"] = image.error

    return result


async def midjourney_list_images(
    limit: int = 10,
    status: str | None = None,
) -> dict[str, Any]:
    """List recently generated images.

    Args:
        limit: Maximum number to retrieve (default: 10, max: 100)
        status: Status to filter by (pending, in-progress, completed, failed)

    Returns:
        Dictionary containing image list info:
        - images: List of image info
        - total: Number of images retrieved

    Raises:
        ImagineAPIError: If API call fails
    """
    limit = min(max(1, limit), 100)

    async with ImagineAPIClient() as client:
        response = await client.list_images(limit=limit, status=status)

    images = []
    for image in response.data:
        item: dict[str, Any] = {
            "id": image.id,
            "prompt": image.prompt[:100] + "..." if len(image.prompt) > 100 else image.prompt,
            "status": image.status.value,
        }

        if image.status == ImageStatus.COMPLETED:
            item["url"] = image.result or image.url

        if image.date_created:
            item["date_created"] = image.date_created

        images.append(item)

    return {
        "images": images,
        "total": len(images),
    }


async def midjourney_get_image(
    image_id: str,
    upscale_index: int | None = None,
    max_width: int = 1024,
) -> Image:
    """Get image data for inline display in Claude Desktop.

    Args:
        image_id: Image ID to retrieve
        upscale_index: Index of upscaled image (0-3). If None, returns grid image.
        max_width: Maximum width for resizing (default: 1024). Images wider than this
            will be resized to fit within MCP's 1MB response limit.

    Returns:
        Image object for inline display

    Raises:
        ValueError: If image is not completed or upscale_index is invalid
        ImagineAPIError: If API call fails or image not found
    """
    async with ImagineAPIClient() as client:
        image = await client.get_image(image_id)

    if image.status != ImageStatus.COMPLETED:
        raise ValueError(
            f"Image '{image_id}' is not completed yet. Current status: {image.status.value}"
        )

    image_url: str
    if upscale_index is not None:
        if not image.upscaled_urls:
            raise ValueError(f"No upscaled images available for '{image_id}'")
        if upscale_index < 0 or upscale_index >= len(image.upscaled_urls):
            raise ValueError(
                f"Invalid upscale index {upscale_index}. "
                f"Valid range: 0-{len(image.upscaled_urls) - 1}"
            )
        image_url = image.upscaled_urls[upscale_index]
    else:
        url = image.result or image.url
        if not url:
            raise ValueError(f"No image URL available for '{image_id}'")
        image_url = url

    async with ImagineAPIClient() as client:
        image_bytes = await client.download_image(image_url)

    resized_bytes = _resize_image(image_bytes, max_width)
    return Image(data=resized_bytes, format="jpeg")


def _resize_image(image_bytes: bytes, max_width: int) -> bytes:
    """Resize image to fit within max_width while maintaining aspect ratio.

    Args:
        image_bytes: Original image bytes
        max_width: Maximum width for the output image

    Returns:
        Resized image as JPEG bytes
    """
    with PILImage.open(io.BytesIO(image_bytes)) as img:
        processed: PILImage.Image = img
        if processed.mode in ("RGBA", "P"):
            processed = processed.convert("RGB")

        if processed.width > max_width:
            ratio = max_width / processed.width
            new_height = int(processed.height * ratio)
            processed = processed.resize((max_width, new_height), PILImage.Resampling.LANCZOS)

        output = io.BytesIO()
        processed.save(output, format="JPEG", quality=85, optimize=True)
        return output.getvalue()
