"""Map Pydantic AI AgentRunResult to Bandito Spans."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from pydantic_ai import AgentRunResult
from pydantic_ai.messages import (
    ModelRequest,
    ModelResponse,
    TextPart,
    ToolCallPart,
    ToolReturnPart,
    UserPromptPart,
)

from ..models import Span


def _response_text(response: ModelResponse) -> str | None:
    texts = [p.content for p in response.parts if isinstance(p, TextPart)]
    return "\n".join(texts) if texts else None


def _tool_calls_by_id(response: ModelResponse) -> dict[str, ToolCallPart]:
    return {
        p.tool_call_id: p
        for p in response.parts
        if isinstance(p, ToolCallPart) and p.tool_call_id
    }


def extract(result: AgentRunResult) -> tuple[Any, list[Span]]:
    """Extract output and spans from a Pydantic AI AgentRunResult."""
    messages = result.all_messages()
    spans: list[Span] = []
    pending_tool_calls: dict[str, ToolCallPart] = {}
    last_llm_span_id: str | None = None

    for msg in messages:
        if isinstance(msg, ModelResponse):
            usage = msg.usage
            token_details: dict[str, int] = {}
            if usage.cache_read_tokens:
                token_details["cache_read"] = usage.cache_read_tokens
            if usage.cache_write_tokens:
                token_details["cache_write"] = usage.cache_write_tokens

            llm_span = Span(
                name=msg.model_name or "unknown",
                kind="llm",
                output=_response_text(msg),
                started_at=msg.timestamp,
                ended_at=msg.timestamp,
                model=msg.model_name,
                provider=msg.provider_name,
                response_model=msg.model_name,
                response_id=msg.provider_response_id,
                input_tokens=usage.input_tokens or None,
                output_tokens=usage.output_tokens or None,
                token_details=token_details or None,
                finish_reason=msg.finish_reason if msg.finish_reason else None,
            )
            spans.append(llm_span)
            last_llm_span_id = llm_span.span_id
            pending_tool_calls.update(_tool_calls_by_id(msg))

        elif isinstance(msg, ModelRequest):
            for part in msg.parts:
                if isinstance(part, ToolReturnPart):
                    call = pending_tool_calls.pop(part.tool_call_id, None) if part.tool_call_id else None
                    input_data = None
                    if call is not None:
                        try:
                            input_data = call.args_as_dict()
                        except Exception:
                            input_data = str(call.args)

                    tool_span = Span(
                        name=part.tool_name,
                        kind="tool",
                        input=input_data,
                        output=part.content,
                        started_at=part.timestamp or datetime.now(timezone.utc),
                        ended_at=part.timestamp or datetime.now(timezone.utc),
                        status="error" if part.outcome == "error" else "ok",
                        parent_id=last_llm_span_id,
                    )
                    spans.append(tool_span)

    return result.output, spans
