"""Task query tools for NanoBanana API."""

from typing import Annotated

from pydantic import Field

from core.client import client
from core.server import mcp
from core.utils import format_task_result


@mcp.tool()
async def nanobanana_get_task(
    task_id: Annotated[
        str,
        Field(
            description="The task ID returned from a generation or edit request. This is the 'task_id' field from any nanobanana_generate_image or nanobanana_edit_image tool response."
        ),
    ],
) -> str:
    """Query the status and result of an image generation or edit task.

    Use this to check if a generation/edit is complete and retrieve the resulting
    image URLs and metadata.

    Use this when:
    - You want to check if an image generation has completed
    - You need to retrieve image URLs from a previous generation
    - You want to get the full details of a generated/edited image

    Returns:
        Task status and image information including URLs and prompts.
    """
    result = await client.query_task(
        id=task_id,
        action="retrieve",
    )
    return format_task_result(result)


@mcp.tool()
async def nanobanana_get_tasks_batch(
    task_ids: Annotated[
        list[str],
        Field(description="List of task IDs to query. Allows querying multiple tasks at once."),
    ],
) -> str:
    """Query multiple image generation/edit tasks at once.

    Efficiently check the status of multiple tasks in a single request.
    More efficient than calling nanobanana_get_task multiple times.

    Use this when:
    - You have multiple pending generations to check
    - You want to get status of several images at once
    - You're tracking a batch of generations

    Returns:
        Status and image information for all queried tasks.
    """
    result = await client.query_task(
        ids=task_ids,
        action="retrieve_batch",
    )

    if "error" in result:
        error = result.get("error", {})
        return f"Error: {error.get('code', 'unknown')} - {error.get('message', 'Unknown error')}"

    lines = [f"Total Tasks: {result.get('count', 0)}", ""]

    for item in result.get("items", []):
        response_info = item.get("response", {})
        lines.extend(
            [
                f"=== Task: {item.get('id', 'N/A')} ===",
                f"Created At: {item.get('created_at', 'N/A')}",
                f"Success: {response_info.get('success', False)}",
            ]
        )

        for image in response_info.get("data", []):
            lines.append(f"  - Image: {image.get('image_url', 'N/A')}")

        lines.append("")

    return "\n".join(lines)
