"""Map Anthropic Message responses to Bandito Spans."""

from __future__ import annotations

from typing import Any

from anthropic.types import Message, TextBlock, ToolUseBlock

from ..models import Span


def extract(message: Message) -> tuple[Any, list[Span]]:
    """Extract output and spans from an Anthropic Message."""
    # Build output text
    text_parts = [b.text for b in message.content if isinstance(b, TextBlock)]
    output = "\n".join(text_parts) if text_parts else None

    # Token details
    usage = message.usage
    token_details: dict[str, int] = {}
    if getattr(usage, "cache_read_input_tokens", None):
        token_details["cache_read"] = usage.cache_read_input_tokens
    if getattr(usage, "cache_creation_input_tokens", None):
        token_details["cache_write"] = usage.cache_creation_input_tokens

    # Map stop_reason to finish_reason
    finish_reason = message.stop_reason  # end_turn, max_tokens, tool_use, etc.

    llm_span = Span(
        name=message.model,
        kind="llm",
        output=output,
        model=message.model,
        provider="anthropic",
        response_model=message.model,
        response_id=message.id,
        input_tokens=usage.input_tokens,
        output_tokens=usage.output_tokens,
        token_details=token_details or None,
        finish_reason=finish_reason,
    )

    spans: list[Span] = [llm_span]

    # Extract tool use blocks as tool spans
    for block in message.content:
        if isinstance(block, ToolUseBlock):
            tool_span = Span(
                name=block.name,
                kind="tool",
                input=block.input,
                parent_id=llm_span.span_id,
            )
            spans.append(tool_span)

    return output, spans
