"""Shared Bedrock Converse API helpers.

These utilities are used by any command that calls the Converse API
(invoke, chat, batch, compare) to avoid duplicating request construction
and response parsing logic.
"""

from __future__ import annotations

from typing import Any

from alice_cli.constants import DEFAULT_MAX_TOKENS
from alice_cli.models import TokenUsage


def build_converse_kwargs(
    model_id: str,
    messages: list[dict[str, Any]],
    system_prompt: str | None = None,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> dict[str, Any]:
    """Build the kwargs dict for a Bedrock Converse API call.

    Args:
        model_id: Resolved model ID or inference profile ARN.
        messages: List of message dicts in Converse API format.
        system_prompt: Optional system prompt text.
        max_tokens: Maximum tokens for the response (default: DEFAULT_MAX_TOKENS).

    Returns:
        A kwargs dict ready to be unpacked into ``client.converse(**kwargs)``.
    """
    kwargs: dict[str, Any] = {
        "modelId": model_id,
        "messages": messages,
        "inferenceConfig": {"maxTokens": max_tokens},
    }
    if system_prompt:
        kwargs["system"] = [{"text": system_prompt}]
    return kwargs


def parse_converse_response(response: dict[str, Any]) -> tuple[str, TokenUsage]:
    """Extract text and token usage from a Converse API response.

    Args:
        response: The raw dict returned by ``client.converse()``.

    Returns:
        A tuple of (response_text, TokenUsage).
        ``response_text`` is the concatenation of all text content blocks.
        ``TokenUsage`` contains input, output, and total token counts.
    """
    output_message = response.get("output", {}).get("message", {})
    content_blocks = output_message.get("content", [])
    text = "\n".join(block["text"] for block in content_blocks if "text" in block)
    usage = response.get("usage", {})
    input_tokens = usage.get("inputTokens", 0)
    output_tokens = usage.get("outputTokens", 0)
    return text, TokenUsage(
        input=input_tokens, output=output_tokens, total=input_tokens + output_tokens
    )
