from __future__ import annotations

import logging
import threading
from collections import OrderedDict
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, Optional, Union
from urllib.parse import urlparse

from agents import MCPListToolsSpanData
from agents.tracing import Span, Trace, TracingProcessor
from agents.tracing.span_data import (
    AgentSpanData,
    CustomSpanData,
    FunctionSpanData,
    GenerationSpanData,
    GuardrailSpanData,
    HandoffSpanData,
    ResponseSpanData,
    SpanData,
)
from openai.types.responses import (
    EasyInputMessageParam,
    FunctionTool,
    Response,
    ResponseCustomToolCall,
    ResponseCustomToolCallOutputParam,
    ResponseCustomToolCallParam,
    ResponseFunctionToolCall,
    ResponseFunctionToolCallParam,
    ResponseInputContentParam,
    ResponseInputItemParam,
    ResponseOutputItem,
    ResponseOutputMessage,
    ResponseOutputMessageParam,
    ResponseOutputRefusal,
    ResponseOutputText,
    ResponseReasoningItem,
    ResponseUsage,
    Tool,
)
from openai.types.responses.response_input_item_param import FunctionCallOutput, Message
from openai.types.responses.response_output_message_param import Content
from opentelemetry.context import attach, detach
from opentelemetry.trace import Span as OtelSpan
from opentelemetry.trace import (
    Status,
    StatusCode,
    Tracer,
    set_span_in_context,
)
from opentelemetry.util.types import AttributeValue
from typing_extensions import assert_never

from openinference.instrumentation import infer_llm_provider_from_host, safe_json_dumps
from openinference.instrumentation.openai_agents._tool_schemas import get_tool_schema
from openinference.semconv.trace import (
    MessageAttributes,
    MessageContentAttributes,
    OpenInferenceLLMSystemValues,
    OpenInferenceMimeTypeValues,
    OpenInferenceSpanKindValues,
    SpanAttributes,
    ToolAttributes,
    ToolCallAttributes,
)

logger = logging.getLogger(__name__)


class OpenInferenceTracingProcessor(TracingProcessor):
    _MAX_HANDOFFS_IN_FLIGHT = 1000

    def __init__(self, tracer: Tracer) -> None:
        self._tracer = tracer
        # Guards the capped insert into _reverse_handoffs_dict below, which is the only
        # read-modify-write here: every other access to the dicts on this object is a
        # single subscript, pop or get. Those are atomic because every key here is a str,
        # so the operation never re-enters the interpreter to hash or compare -- a key
        # with a Python-level __hash__ or __eq__ would not get that. The SDK invokes these
        # callbacks on whichever thread is running the agent, so a process driving
        # Runner.run_sync from several threads reaches them concurrently.
        self._lock = threading.Lock()
        # Spans and traces still in flight, deliberately unbounded. The number of
        # simultaneously open spans scales with how many runs are in flight, so any cap
        # large enough to be safe under load is too large to bound anything, and evicting
        # an open span is far more damaging than leaking a dict entry: the span would be
        # ended early with only its start attributes, its real end would be dropped, and
        # its children would reparent to the trace root. Entries for spans that never end
        # (a hard cancellation) are therefore accepted as a leak rather than reaped.
        self._root_spans: dict[str, OtelSpan] = {}
        self._otel_spans: dict[str, OtelSpan] = {}
        self._tokens: dict[str, object] = {}
        # This captures in flight handoff. Once the handoff is complete, the entry is deleted
        # If the handoff does not complete, the entry stays in the dict.
        # Use an OrderedDict and _MAX_HANDOFFS_IN_FLIGHT to cap the size of the dict
        # in case there are large numbers of orphaned handoffs
        self._reverse_handoffs_dict: OrderedDict[str, str] = OrderedDict()

    def on_trace_start(self, trace: Trace) -> None:
        """Called when a trace is started.

        Args:
            trace: The trace that started.
        """
        otel_span = self._tracer.start_span(
            name=trace.name,
            attributes={
                OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.AGENT.value,
            },
        )
        self._root_spans[trace.trace_id] = otel_span

    def on_trace_end(self, trace: Trace) -> None:
        """Called when a trace is finished.

        Args:
            trace: The trace that started.
        """
        root_span = self._root_spans.pop(trace.trace_id, None)
        if root_span:
            root_span.set_status(Status(StatusCode.OK))
            root_span.end()

    def on_span_start(self, span: Span[Any]) -> None:
        """Called when a span is started.

        Args:
            span: The span that started.
        """
        if not span.started_at:
            return
        start_time = datetime.fromisoformat(span.started_at)
        parent_span = (
            self._otel_spans.get(span.parent_id)
            if span.parent_id
            else self._root_spans.get(span.trace_id)
        )
        context = set_span_in_context(parent_span) if parent_span else None
        span_name = _get_span_name(span)
        span_kind = _get_span_kind(span.span_data)
        attributes: dict[str, AttributeValue] = {OPENINFERENCE_SPAN_KIND: span_kind}
        # llm.system describes the LLM being called, so it belongs on LLM spans only
        # rather than on every agent, tool, and handoff span in the trace.
        if span_kind == OpenInferenceSpanKindValues.LLM.value:
            attributes[LLM_SYSTEM] = OpenInferenceLLMSystemValues.OPENAI.value
        otel_span = self._tracer.start_span(
            name=span_name,
            context=context,
            start_time=_as_utc_nano(start_time),
            attributes=attributes,
        )
        self._otel_spans[span.span_id] = otel_span
        self._tokens[span.span_id] = attach(set_span_in_context(otel_span))

    def on_span_end(self, span: Span[Any]) -> None:
        """Called when a span is finished. Should not block or raise exceptions.

        Args:
            span: The span that finished.
        """
        token = self._tokens.pop(span.span_id, None)
        otel_span = self._otel_spans.pop(span.span_id, None)
        if token:
            detach(token)  # type: ignore[arg-type]
        if otel_span is None:
            return
        otel_span.update_name(_get_span_name(span))
        data = span.span_data
        if isinstance(data, ResponseSpanData):
            if isinstance(response := data.response, Response):
                otel_span.set_attribute(OUTPUT_MIME_TYPE, JSON)
                otel_span.set_attribute(OUTPUT_VALUE, response.model_dump_json())
                for k, v in _get_attributes_from_response(response):
                    otel_span.set_attribute(k, v)
            if input := data.input:
                if isinstance(input, str):
                    otel_span.set_attribute(INPUT_VALUE, input)
                elif isinstance(input, list):
                    otel_span.set_attribute(INPUT_MIME_TYPE, JSON)
                    otel_span.set_attribute(INPUT_VALUE, safe_json_dumps(input))
                    for k, v in _get_attributes_from_input(input):
                        otel_span.set_attribute(k, v)
                elif TYPE_CHECKING:
                    assert_never(input)
        elif isinstance(data, GenerationSpanData):
            for k, v in _get_attributes_from_generation_span_data(data):
                otel_span.set_attribute(k, v)
        elif isinstance(data, FunctionSpanData):
            for k, v in _get_attributes_from_function_span_data(data):
                otel_span.set_attribute(k, v)
            # Function spans carry no schema of their own, so read it off the live tool
            # object being invoked. See _tool_schemas for how that is made available.
            if (schema := get_tool_schema(data.name)) is not None:
                description, parameters = schema
                if description is not None:
                    otel_span.set_attribute(TOOL_DESCRIPTION, description)
                if parameters is not None:
                    otel_span.set_attribute(TOOL_PARAMETERS, parameters)
        elif isinstance(data, MCPListToolsSpanData):
            for k, v in _get_attributes_from_mcp_list_tool_span_data(data):
                otel_span.set_attribute(k, v)
        elif isinstance(data, HandoffSpanData):
            # Handoffs surface as TOOL spans but the SDK leaves them without input or
            # output, so record the agents the handoff moved between. Note that
            # HandoffSpanData carries only these names -- the name of the tool the model
            # called is deliberately not reconstructed here, since a handoff created
            # with tool_name_override would not match a reconstructed value.
            if data.from_agent is not None:
                otel_span.set_attribute(INPUT_VALUE, data.from_agent)
            if data.to_agent is not None:
                otel_span.set_attribute(OUTPUT_VALUE, data.to_agent)
            # Set this dict to find the parent node when the agent span starts
            if data.to_agent and data.from_agent:
                key = f"{data.to_agent}:{span.trace_id}"
                with self._lock:
                    self._reverse_handoffs_dict[key] = data.from_agent
                    # Cap the size of the dict
                    while len(self._reverse_handoffs_dict) > self._MAX_HANDOFFS_IN_FLIGHT:
                        self._reverse_handoffs_dict.popitem(last=False)
        elif isinstance(data, AgentSpanData):
            otel_span.set_attribute(GRAPH_NODE_ID, data.name)
            otel_span.set_attribute(AGENT_NAME, data.name)
            # Lookup the parent node if exists
            key = f"{data.name}:{span.trace_id}"
            parent_node = self._reverse_handoffs_dict.pop(key, None)
            if parent_node:
                otel_span.set_attribute(GRAPH_NODE_PARENT_ID, parent_node)

        end_time: Optional[int] = None
        if span.ended_at:
            try:
                end_time = _as_utc_nano(datetime.fromisoformat(span.ended_at))
            except ValueError:
                pass
        otel_span.set_status(status=_get_span_status(span))
        otel_span.end(end_time)

    def force_flush(self) -> None:
        """Forces an immediate flush of all queued spans/traces."""
        # TODO
        pass

    def shutdown(self) -> None:
        """Called when the application stops."""
        # TODO
        pass


def _as_utc_nano(dt: datetime) -> int:
    return int(dt.astimezone(timezone.utc).timestamp() * 1_000_000_000)


def _get_span_name(obj: Span[Any]) -> str:
    if hasattr(data := obj.span_data, "name") and isinstance(name := data.name, str):
        return name
    if isinstance(obj.span_data, HandoffSpanData) and obj.span_data.to_agent:
        return f"handoff to {obj.span_data.to_agent}"
    return obj.span_data.type  # type: ignore[no-any-return]


def _get_span_kind(obj: SpanData) -> str:
    if isinstance(obj, AgentSpanData):
        return OpenInferenceSpanKindValues.AGENT.value
    if isinstance(obj, FunctionSpanData):
        return OpenInferenceSpanKindValues.TOOL.value
    if isinstance(obj, GenerationSpanData):
        return OpenInferenceSpanKindValues.LLM.value
    if isinstance(obj, ResponseSpanData):
        return OpenInferenceSpanKindValues.LLM.value
    if isinstance(obj, HandoffSpanData):
        return OpenInferenceSpanKindValues.TOOL.value
    if isinstance(obj, CustomSpanData):
        return OpenInferenceSpanKindValues.CHAIN.value
    if isinstance(obj, GuardrailSpanData):
        return OpenInferenceSpanKindValues.GUARDRAIL.value
    return OpenInferenceSpanKindValues.CHAIN.value


def _get_attributes_from_input(
    obj: Iterable[ResponseInputItemParam],
    msg_idx: int = 1,
) -> Iterator[tuple[str, AttributeValue]]:
    for i, item in enumerate(obj, msg_idx):
        prefix = f"{LLM_INPUT_MESSAGES}.{i}."
        if "type" not in item:
            if "role" in item and "content" in item:
                yield from _get_attributes_from_message_param(
                    {  # type: ignore[misc, arg-type]
                        "type": "message",
                        "role": item["role"],  # type: ignore[typeddict-item]
                        "content": item["content"],  # type: ignore[typeddict-item]
                    },
                    prefix,
                )
        elif item["type"] == "message":
            yield from _get_attributes_from_message_param(item, prefix)
        elif item["type"] == "file_search_call":
            continue  # TODO
        elif item["type"] == "computer_call":
            continue  # TODO
        elif item["type"] == "computer_call_output":
            continue  # TODO
        elif item["type"] == "web_search_call":
            continue  # TODO
        elif item["type"] == "function_call":
            yield f"{prefix}{MESSAGE_ROLE}", "assistant"
            yield from _get_attributes_from_response_function_tool_call_param(
                item,
                f"{prefix}{MESSAGE_TOOL_CALLS}.0.",
            )
        elif item["type"] == "function_call_output":
            yield from _get_attributes_from_function_call_output(item, prefix)
        elif item["type"] == "custom_tool_call":
            yield f"{prefix}{MessageAttributes.MESSAGE_ROLE}", "assistant"
            yield from _get_attributes_from_response_custom_tool_call_param(
                item,
                f"{prefix}{MessageAttributes.MESSAGE_TOOL_CALLS}.0.",
            )
        elif item["type"] == "custom_tool_call_output":
            yield from _get_attributes_from_response_custom_tool_call_output_param(item, prefix)
        elif item["type"] == "reasoning":
            continue  # TODO
        elif item["type"] == "item_reference":
            continue  # TODO
        elif item["type"] == "image_generation_call":
            continue  # TODO
        elif item["type"] == "code_interpreter_call":
            continue  # TODO
        elif item["type"] == "local_shell_call":
            continue  # TODO
        elif item["type"] == "local_shell_call_output":
            continue  # TODO
        elif item["type"] == "mcp_list_tools":
            continue  # TODO
        elif item["type"] == "mcp_approval_request":
            continue  # TODO
        elif item["type"] == "mcp_approval_response":
            continue  # TODO
        elif item["type"] == "mcp_call":
            continue  # TODO
        elif item["type"] == "shell_call":
            # TODO: Handle shell call
            continue
        elif item["type"] == "shell_call_output":
            # TODO: Handle shell call output
            continue
        elif item["type"] == "apply_patch_call":
            # TODO: Handle patch call
            continue
        elif item["type"] == "apply_patch_call_output":
            # TODO: Handle patch call output
            continue
        elif item["type"] == "compaction":
            # TODO: Handle compaction response
            continue
        elif item["type"] == "tool_search_call":
            # TODO: Handle tool search call response
            continue
        elif item["type"] == "tool_search_output":
            # TODO: Handle tool search output response
            continue
        elif item["type"] == "compaction_trigger":
            # TODO: Handle compaction trigger response
            continue
        elif item["type"] == "additional_tools":
            # TODO: Handle additional tools
            continue
        elif item["type"] == "program":
            # TODO: Handle program
            continue
        elif item["type"] == "program_output":
            # TODO: Handle program output
            continue
        elif TYPE_CHECKING and item["type"] is not None:
            assert_never(item["type"])


def _get_attributes_from_message_param(
    obj: Union[
        EasyInputMessageParam,
        Message,
        ResponseOutputMessageParam,
    ],
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    yield f"{prefix}{MESSAGE_ROLE}", obj["role"]
    if content := obj.get("content"):
        if isinstance(content, str):
            yield f"{prefix}{MESSAGE_CONTENT}", content
        elif isinstance(content, list):
            yield from _get_attributes_from_message_content_list(content, prefix)


def _get_attributes_from_response_function_tool_call_param(
    obj: ResponseFunctionToolCallParam,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    yield f"{prefix}{TOOL_CALL_ID}", obj["call_id"]
    yield f"{prefix}{TOOL_CALL_FUNCTION_NAME}", obj["name"]
    if obj["arguments"] != "{}":
        yield f"{prefix}{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}", obj["arguments"]


def _get_attributes_from_response_custom_tool_call_param(
    obj: ResponseCustomToolCallParam,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    if (call_id := obj.get("call_id")) is not None:
        yield f"{prefix}{ToolCallAttributes.TOOL_CALL_ID}", call_id
    if (name := obj.get("name")) is not None:
        yield f"{prefix}{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}", name
    if (input_data := obj.get("input")) is not None:
        yield (
            f"{prefix}{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
            safe_json_dumps({"input": input_data}),
        )


def _get_attributes_from_response_custom_tool_call_output_param(
    obj: ResponseCustomToolCallOutputParam,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    yield f"{prefix}{MessageAttributes.MESSAGE_ROLE}", "tool"
    if (call_id := obj.get("call_id")) is not None:
        yield f"{prefix}{MessageAttributes.MESSAGE_TOOL_CALL_ID}", call_id
    if "output" in obj:
        output = obj["output"]
        if output is not None:
            if isinstance(output, str):
                output_value = output
            else:
                output_value = safe_json_dumps(output)
            yield f"{prefix}{MessageAttributes.MESSAGE_CONTENT}", output_value


def _get_attributes_from_function_call_output(
    obj: FunctionCallOutput,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    yield f"{prefix}{MESSAGE_ROLE}", "tool"
    yield f"{prefix}{MESSAGE_TOOL_CALL_ID}", obj["call_id"]
    output = obj["output"]
    if output is not None:
        if isinstance(output, str):
            output_value = output
        else:
            output_value = safe_json_dumps(output)
        yield f"{prefix}{MESSAGE_CONTENT}", output_value


def _get_attributes_from_generation_span_data(
    obj: GenerationSpanData,
) -> Iterator[tuple[str, AttributeValue]]:
    if isinstance(model := obj.model, str):
        yield LLM_MODEL_NAME, model
    if isinstance(obj.model_config, dict) and (
        param := {k: v for k, v in obj.model_config.items() if v is not None}
    ):
        yield LLM_INVOCATION_PARAMETERS, safe_json_dumps(param)
        if isinstance(base_url := param.get("base_url"), str) and (
            host := urlparse(base_url).hostname
        ):
            if provider := infer_llm_provider_from_host(host):
                yield LLM_PROVIDER, provider.value
    yield from _get_attributes_from_chat_completions_input(obj.input)
    yield from _get_attributes_from_chat_completions_output(obj.output)
    yield from _get_attributes_from_chat_completions_usage(obj.usage)


def _get_attributes_from_mcp_list_tool_span_data(
    obj: MCPListToolsSpanData,
) -> Iterator[tuple[str, AttributeValue]]:
    yield OUTPUT_VALUE, safe_json_dumps(obj.result)
    yield OUTPUT_MIME_TYPE, JSON


def _get_attributes_from_chat_completions_input(
    obj: Optional[Iterable[Mapping[str, Any]]],
) -> Iterator[tuple[str, AttributeValue]]:
    if not obj:
        return
    try:
        yield INPUT_VALUE, safe_json_dumps(obj)
        yield INPUT_MIME_TYPE, JSON
    except Exception:
        pass
    yield from _get_attributes_from_chat_completions_message_dicts(
        obj,
        f"{LLM_INPUT_MESSAGES}.",
    )


def _get_attributes_from_chat_completions_output(
    obj: Optional[Iterable[Mapping[str, Any]]],
) -> Iterator[tuple[str, AttributeValue]]:
    if not obj:
        return
    try:
        yield OUTPUT_VALUE, safe_json_dumps(obj)
        yield OUTPUT_MIME_TYPE, JSON
    except Exception:
        pass
    yield from _get_attributes_from_chat_completions_message_dicts(
        obj,
        f"{LLM_OUTPUT_MESSAGES}.",
    )


def _get_attributes_from_chat_completions_message_dicts(
    obj: Iterable[Mapping[str, Any]],
    prefix: str = "",
    msg_idx: int = 0,
    tool_call_idx: int = 0,
) -> Iterator[tuple[str, AttributeValue]]:
    if not isinstance(obj, Iterable):
        return
    for msg in obj:
        if isinstance(role := msg.get("role"), str):
            yield f"{prefix}{msg_idx}.{MESSAGE_ROLE}", role
        if content := msg.get("content"):
            yield from _get_attributes_from_chat_completions_message_content(
                content,
                f"{prefix}{msg_idx}.",
            )
        if isinstance(tool_call_id := msg.get("tool_call_id"), str):
            yield f"{prefix}{msg_idx}.{MESSAGE_TOOL_CALL_ID}", tool_call_id
        if isinstance(tool_calls := msg.get("tool_calls"), Iterable):
            for tc in tool_calls:
                yield from _get_attributes_from_chat_completions_tool_call_dict(
                    tc,
                    f"{prefix}{msg_idx}.{MESSAGE_TOOL_CALLS}.{tool_call_idx}.",
                )
                tool_call_idx += 1
        msg_idx += 1


def _get_attributes_from_chat_completions_message_content(
    obj: Union[str, Iterable[Mapping[str, Any]]],
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    if isinstance(obj, str):
        yield f"{prefix}{MESSAGE_CONTENT}", obj
    elif isinstance(obj, Iterable):
        for i, item in enumerate(obj):
            if not isinstance(item, Mapping):
                continue
            yield from _get_attributes_from_chat_completions_message_content_item(
                item,
                f"{prefix}{MESSAGE_CONTENTS}.{i}.",
            )


def _get_attributes_from_chat_completions_message_content_item(
    obj: Mapping[str, Any],
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    if obj.get("type") == "text" and (text := obj.get("text")):
        yield f"{prefix}{MESSAGE_CONTENT_TYPE}", "text"
        yield f"{prefix}{MESSAGE_CONTENT_TEXT}", text


def _get_attributes_from_chat_completions_tool_call_dict(
    obj: Mapping[str, Any],
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    if id_ := obj.get("id"):
        yield f"{prefix}{TOOL_CALL_ID}", id_
    if function := obj.get("function"):
        if name := function.get("name"):
            yield f"{prefix}{TOOL_CALL_FUNCTION_NAME}", name
        if arguments := function.get("arguments"):
            if arguments != "{}":
                yield f"{prefix}{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}", arguments


def _get_attributes_from_chat_completions_usage(
    obj: Optional[Mapping[str, Any]],
) -> Iterator[tuple[str, AttributeValue]]:
    if not obj:
        return
    if input_tokens := obj.get("input_tokens"):
        yield LLM_TOKEN_COUNT_PROMPT, input_tokens
    if output_tokens := obj.get("output_tokens"):
        yield LLM_TOKEN_COUNT_COMPLETION, output_tokens


# convert dict, tuple, etc into one of these types ['bool', 'str', 'bytes', 'int', 'float']
def _convert_to_primitive(value: Any) -> Union[bool, str, bytes, int, float]:
    if isinstance(value, (bool, str, bytes, int, float)):
        return value
    if isinstance(value, (list, tuple)):
        return safe_json_dumps(value)
    if isinstance(value, dict):
        return safe_json_dumps(value)
    return str(value)


def _get_attributes_from_function_span_data(
    obj: FunctionSpanData,
) -> Iterator[tuple[str, AttributeValue]]:
    yield TOOL_NAME, obj.name
    if obj.input:
        yield INPUT_VALUE, obj.input
        yield INPUT_MIME_TYPE, JSON
    if obj.output is not None:
        yield OUTPUT_VALUE, _convert_to_primitive(obj.output)
        if (
            isinstance(obj.output, str)
            and len(obj.output) > 1
            and obj.output[0] == "{"
            and obj.output[-1] == "}"
        ):
            yield OUTPUT_MIME_TYPE, JSON


def _get_attributes_from_message_content_list(
    obj: Iterable[Union[ResponseInputContentParam, Content]],
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    for i, item in enumerate(obj):
        if item["type"] == "input_text" or item["type"] == "output_text":
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TYPE}", "text"
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TEXT}", item["text"]
        elif item["type"] == "input_image":
            # TODO
            ...
        elif item["type"] == "input_file":
            # TODO
            ...
        elif item["type"] == "refusal":
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TYPE}", "text"
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TEXT}", item["refusal"]
        elif item["type"] == "input_audio":
            # TODO: Handle input audio (OpenAI 1.105.0+)
            ...
        elif TYPE_CHECKING:
            assert_never(item["type"])


def _get_attributes_from_response(obj: Response) -> Iterator[tuple[str, AttributeValue]]:
    yield from _get_attributes_from_tools(obj.tools)
    yield from _get_attributes_from_usage(obj.usage)
    yield from _get_attributes_from_response_output(obj.output)
    if isinstance(obj.instructions, str):
        yield from _get_attributes_from_response_instruction(obj.instructions)
    else:
        pass  # TODO: handle list instructions
    yield LLM_MODEL_NAME, obj.model
    param = obj.model_dump(
        exclude_none=True,
        exclude={"object", "tools", "usage", "output", "error", "status"},
    )
    yield LLM_INVOCATION_PARAMETERS, safe_json_dumps(param)


def _get_attributes_from_tools(
    tools: Optional[Iterable[Tool]],
) -> Iterator[tuple[str, AttributeValue]]:
    if not tools:
        return
    for i, tool in enumerate(tools):
        if isinstance(tool, FunctionTool):
            yield (
                f"{LLM_TOOLS}.{i}.{TOOL_JSON_SCHEMA}",
                safe_json_dumps(
                    {
                        "type": "function",
                        "function": {
                            "name": tool.name,
                            "description": tool.description,
                            "parameters": tool.parameters,
                            "strict": tool.strict,
                        },
                    }
                ),
            )
        else:
            pass


def _get_attributes_from_response_output(
    obj: Iterable[ResponseOutputItem],
    msg_idx: int = 0,
) -> Iterator[tuple[str, AttributeValue]]:
    tool_call_idx = 0
    for item in obj:
        if item.type == "message":
            prefix = f"{LLM_OUTPUT_MESSAGES}.{msg_idx}."
            yield from _get_attributes_from_message(item, prefix)
            msg_idx += 1
        elif item.type == "function_call":
            yield f"{LLM_OUTPUT_MESSAGES}.{msg_idx}.{MESSAGE_ROLE}", "assistant"
            prefix = f"{LLM_OUTPUT_MESSAGES}.{msg_idx}.{MESSAGE_TOOL_CALLS}.{tool_call_idx}."
            yield from _get_attributes_from_function_tool_call(item, prefix)
            tool_call_idx += 1
        elif item.type == "custom_tool_call":
            yield f"{LLM_OUTPUT_MESSAGES}.{msg_idx}.{MESSAGE_ROLE}", "assistant"
            prefix = f"{LLM_OUTPUT_MESSAGES}.{msg_idx}.{MESSAGE_TOOL_CALLS}.{tool_call_idx}."
            yield from _get_attributes_from_response_custom_tool_call(item, prefix)
            tool_call_idx += 1
        elif item.type == "file_search_call":
            ...  # TODO
        elif item.type == "web_search_call":
            ...  # TODO
        elif item.type == "computer_call":
            ...  # TODO
        elif item.type == "reasoning":
            prefix = f"{LLM_OUTPUT_MESSAGES}.{msg_idx}."
            attrs = list(_get_attributes_from_reasoning_item(item, prefix))
            if attrs:
                for k, v in attrs:
                    yield k, v
                msg_idx += 1
        elif item.type == "image_generation_call":
            ...  # TODO
        elif item.type == "code_interpreter_call":
            ...  # TODO
        elif item.type == "local_shell_call":
            ...  # TODO
        elif item.type == "mcp_call":
            ...  # TODO
        elif item.type == "mcp_list_tools":
            ...  # TODO
        elif item.type == "mcp_approval_request":
            ...  # TODO
        elif item.type == "shell_call":
            ...  # TODO
        elif item.type == "shell_call_output":
            ...  # TODO
        elif item.type == "apply_patch_call":
            ...  # TODO
        elif item.type == "apply_patch_call_output":
            ...  # TODO
        elif item.type == "compaction":
            ...  # TODO
        elif item.type == "tool_search_call":
            ...  # TODO
        elif item.type == "tool_search_output":
            ...  # TODO
        elif item.type == "function_call_output":
            ...  # TODO
        elif item.type == "computer_call_output":
            ...  # TODO
        elif item.type == "local_shell_call_output":
            ...  # TODO
        elif item.type == "mcp_approval_response":
            ...  # TODO
        elif item.type == "custom_tool_call_output":
            ...  # TODO
        elif item.type == "additional_tools":
            ...  # TODO
        elif item.type == "program":
            ...  # TODO
        elif item.type == "program_output":
            ...  # TODO
        elif TYPE_CHECKING:
            assert_never(item)


def _get_attributes_from_response_instruction(
    instructions: Optional[str],
) -> Iterator[tuple[str, AttributeValue]]:
    if not instructions:
        return
    yield f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_ROLE}", "system"
    yield f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_CONTENT}", instructions


def _get_attributes_from_function_tool_call(
    obj: ResponseFunctionToolCall,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    yield f"{prefix}{TOOL_CALL_ID}", obj.call_id
    yield f"{prefix}{TOOL_CALL_FUNCTION_NAME}", obj.name
    if obj.arguments != "{}":
        yield f"{prefix}{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}", obj.arguments


def _get_attributes_from_response_custom_tool_call(
    obj: ResponseCustomToolCall,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    if (call_id := obj.call_id) is not None:
        yield f"{prefix}{ToolCallAttributes.TOOL_CALL_ID}", call_id
    if (name := obj.name) is not None:
        yield f"{prefix}{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}", name
    if (input_data := obj.input) is not None:
        yield (
            f"{prefix}{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
            safe_json_dumps({"input": input_data}),
        )


def _get_attributes_from_message(
    obj: ResponseOutputMessage,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    yield f"{prefix}{MESSAGE_ROLE}", obj.role
    text_parts: list[str] = []
    for i, item in enumerate(obj.content):
        if isinstance(item, ResponseOutputText):
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TYPE}", "text"
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TEXT}", item.text
            text_parts.append(item.text)
        elif isinstance(item, ResponseOutputRefusal):
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TYPE}", "text"
            yield f"{prefix}{MESSAGE_CONTENTS}.{i}.{MESSAGE_CONTENT_TEXT}", item.refusal
            text_parts.append(item.refusal)
        elif TYPE_CHECKING:
            assert_never(item)
    # Emit a flat message.content for consumers that don't read the multi-part form.
    if text_parts:
        yield f"{prefix}{MESSAGE_CONTENT}", "\n".join(text_parts)


def _get_attributes_from_usage(
    obj: Optional[ResponseUsage],
) -> Iterator[tuple[str, AttributeValue]]:
    if not obj:
        return
    yield LLM_TOKEN_COUNT_COMPLETION, obj.output_tokens
    yield LLM_TOKEN_COUNT_PROMPT, obj.input_tokens
    yield LLM_TOKEN_COUNT_TOTAL, obj.total_tokens
    if obj.input_tokens_details:
        yield LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, obj.input_tokens_details.cached_tokens
    if obj.output_tokens_details:
        yield (
            LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING,
            obj.output_tokens_details.reasoning_tokens,
        )


def _get_attributes_from_reasoning_item(
    obj: ResponseReasoningItem,
    prefix: str = "",
) -> Iterator[tuple[str, AttributeValue]]:
    summary_texts = [s.text for s in obj.summary if getattr(s, "text", None)]
    content_texts = [c.text for c in (obj.content or []) if getattr(c, "text", None)]
    texts = summary_texts or content_texts
    if not texts and not obj.encrypted_content:
        return

    yield f"{prefix}{MESSAGE_ROLE}", "assistant"
    content_prefix = f"{prefix}{MESSAGE_CONTENTS}.0."
    yield f"{content_prefix}{MESSAGE_CONTENT_TYPE}", "reasoning"
    if texts:
        yield f"{content_prefix}{MESSAGE_CONTENT_TEXT}", "\n\n".join(texts)
    if obj.encrypted_content:
        yield f"{content_prefix}{MESSAGE_CONTENT_ENCRYPTED_CONTENT}", obj.encrypted_content
    if obj.id:
        yield f"{content_prefix}{MESSAGE_CONTENT_ID}", obj.id


def _get_span_status(obj: Span[Any]) -> Status:
    if error := getattr(obj, "error", None):
        return Status(
            status_code=StatusCode.ERROR, description=f"{error.get('message')}: {error.get('data')}"
        )
    else:
        return Status(StatusCode.OK)


INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
INPUT_VALUE = SpanAttributes.INPUT_VALUE
LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
LLM_INVOCATION_PARAMETERS = SpanAttributes.LLM_INVOCATION_PARAMETERS
LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
LLM_PROVIDER = SpanAttributes.LLM_PROVIDER
LLM_SYSTEM = SpanAttributes.LLM_SYSTEM
LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ
LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = (
    SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING
)
LLM_TOOLS = SpanAttributes.LLM_TOOLS
METADATA = SpanAttributes.METADATA
OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
TOOL_DESCRIPTION = SpanAttributes.TOOL_DESCRIPTION
TOOL_NAME = SpanAttributes.TOOL_NAME
TOOL_PARAMETERS = SpanAttributes.TOOL_PARAMETERS
GRAPH_NODE_ID = SpanAttributes.GRAPH_NODE_ID
GRAPH_NODE_PARENT_ID = SpanAttributes.GRAPH_NODE_PARENT_ID
AGENT_NAME = SpanAttributes.AGENT_NAME

MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
MESSAGE_CONTENTS = MessageAttributes.MESSAGE_CONTENTS
MESSAGE_CONTENT_ID = MessageContentAttributes.MESSAGE_CONTENT_ID
MESSAGE_CONTENT_ENCRYPTED_CONTENT = MessageContentAttributes.MESSAGE_CONTENT_ENCRYPTED_CONTENT
MESSAGE_CONTENT_IMAGE = MessageContentAttributes.MESSAGE_CONTENT_IMAGE
MESSAGE_CONTENT_TEXT = MessageContentAttributes.MESSAGE_CONTENT_TEXT
MESSAGE_CONTENT_TYPE = MessageContentAttributes.MESSAGE_CONTENT_TYPE
MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON = MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON
MESSAGE_FUNCTION_CALL_NAME = MessageAttributes.MESSAGE_FUNCTION_CALL_NAME
MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS
MESSAGE_TOOL_CALL_ID = MessageAttributes.MESSAGE_TOOL_CALL_ID

TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME
TOOL_CALL_ID = ToolCallAttributes.TOOL_CALL_ID

TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA

JSON = OpenInferenceMimeTypeValues.JSON.value
