"""Map OpenAI ChatCompletion responses to Bandito Spans.

Also handles OpenRouter and any OpenAI-compatible API that returns ChatCompletion.
"""

from __future__ import annotations

import json
from typing import Any

from openai.types.chat import ChatCompletion

from ..models import Span


def extract(completion: ChatCompletion) -> tuple[Any, list[Span]]:
    """Extract output and spans from an OpenAI ChatCompletion."""
    choice = completion.choices[0] if completion.choices else None
    message = choice.message if choice else None

    output = message.content if message else None
    finish_reason = choice.finish_reason if choice else None

    # Token usage
    usage = completion.usage
    input_tokens = usage.prompt_tokens if usage else None
    output_tokens = usage.completion_tokens if usage else None

    token_details: dict[str, int] = {}
    if usage:
        if usage.prompt_tokens_details:
            details = usage.prompt_tokens_details
            if getattr(details, "cached_tokens", None):
                token_details["cache_read"] = details.cached_tokens
        if usage.completion_tokens_details:
            details = usage.completion_tokens_details
            if getattr(details, "reasoning_tokens", None):
                token_details["reasoning"] = details.reasoning_tokens

    # Provider data
    provider_data: dict[str, Any] = {}
    if completion.system_fingerprint:
        provider_data["system_fingerprint"] = completion.system_fingerprint
    if completion.service_tier:
        provider_data["service_tier"] = completion.service_tier

    llm_span = Span(
        name=completion.model,
        kind="llm",
        output=output,
        model=completion.model,
        provider="openai",
        response_model=completion.model,
        response_id=completion.id,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        token_details=token_details or None,
        finish_reason=finish_reason,
        provider_data=provider_data or None,
    )

    spans: list[Span] = [llm_span]

    # Extract tool calls as tool spans
    if message and message.tool_calls:
        for tool_call in message.tool_calls:
            input_data = None
            if tool_call.function.arguments:
                try:
                    input_data = json.loads(tool_call.function.arguments)
                except (json.JSONDecodeError, TypeError):
                    input_data = tool_call.function.arguments

            tool_span = Span(
                name=tool_call.function.name,
                kind="tool",
                input=input_data,
                parent_id=llm_span.span_id,
                metadata={"tool_call_id": tool_call.id} if tool_call.id else None,
            )
            spans.append(tool_span)

    return output, spans
