"""Map LangChain AIMessage responses to Bandito Spans."""

from __future__ import annotations

from typing import Any

from langchain_core.messages import AIMessage

from ..models import Span


def extract(message: AIMessage) -> tuple[Any, list[Span]]:
    """Extract output and spans from a LangChain AIMessage."""
    output = message.content if isinstance(message.content, str) else str(message.content)

    # Token usage from usage_metadata (standardized across providers)
    usage = message.usage_metadata
    input_tokens = None
    output_tokens = None
    token_details: dict[str, int] = {}
    if usage:
        input_tokens = usage.get("input_tokens")
        output_tokens = usage.get("output_tokens")
        if usage.get("input_token_details"):
            for k, v in usage["input_token_details"].items():
                if v:
                    token_details[k] = v
        if usage.get("output_token_details"):
            for k, v in usage["output_token_details"].items():
                if v:
                    token_details[k] = v

    # Model info from response_metadata
    resp_meta = message.response_metadata or {}
    model = resp_meta.get("model_name") or resp_meta.get("model")
    finish_reason = resp_meta.get("finish_reason") or resp_meta.get("stop_reason")

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

    spans: list[Span] = [llm_span]

    # Extract tool calls
    if message.tool_calls:
        for tool_call in message.tool_calls:
            tool_span = Span(
                name=tool_call.get("name", "unknown"),
                kind="tool",
                input=tool_call.get("args"),
                parent_id=llm_span.span_id,
                metadata={"tool_call_id": tool_call["id"]} if tool_call.get("id") else None,
            )
            spans.append(tool_span)

    return output, spans
