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

from typing import Any

from mcp.server.fastmcp.utilities.types import Image

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,
) -> 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.

    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)

    return Image(data=image_bytes, format="png")
