"""
Implements logging integration with Datadog's LLM Observability Service


API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards

"""

import asyncio
import json
import os
from litellm._uuid import uuid
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Union

import httpx

import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_mock_client import (
    should_use_datadog_mock,
    create_mock_datadog_client,
)
from litellm.integrations.datadog.datadog_handler import (
    get_datadog_service,
    get_datadog_tags,
    get_datadog_base_url_from_env,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.prompt_templates.common_utils import (
    handle_any_messages_to_chat_completion_str_messages_conversion,
)
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.types.integrations.datadog_llm_obs import *
from litellm.types.utils import (
    CallTypes,
    StandardLoggingGuardrailInformation,
    StandardLoggingPayload,
    StandardLoggingPayloadErrorInformation,
)


class DataDogLLMObsLogger(CustomBatchLogger):
    def __init__(self, **kwargs):
        try:
            verbose_logger.debug("DataDogLLMObs: Initializing logger")

            self.is_mock_mode = should_use_datadog_mock()

            if self.is_mock_mode:
                create_mock_datadog_client()
                verbose_logger.debug(
                    "[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode"
                )

            # Configure DataDog endpoint (Agent or Direct API)
            # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
            # Check for agent mode FIRST - agent mode doesn't require DD_API_KEY or DD_SITE
            dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")

            self.async_client = get_async_httpx_client(
                llm_provider=httpxSpecialProvider.LoggingCallback
            )
            self.DD_API_KEY = os.getenv("DD_API_KEY")

            if dd_agent_host:
                self._configure_dd_agent(dd_agent_host=dd_agent_host)
            else:
                # Only require DD_API_KEY and DD_SITE for direct API mode
                if os.getenv("DD_API_KEY", None) is None:
                    raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'")
                if os.getenv("DD_SITE", None) is None:
                    raise Exception(
                        "DD_SITE is not set, set 'DD_SITE=<>', example sit = `us5.datadoghq.com`"
                    )
                self._configure_dd_direct_api()

            # Optional override for testing
            dd_base_url = get_datadog_base_url_from_env()
            if dd_base_url:
                self.intake_url = f"{dd_base_url}/api/intake/llm-obs/v1/trace/spans"

            asyncio.create_task(self.periodic_flush())
            self.flush_lock = asyncio.Lock()
            self.log_queue: List[LLMObsPayload] = []

            #########################################################
            # Handle datadog_llm_observability_params set as litellm.datadog_llm_observability_params
            #########################################################
            dict_datadog_llm_obs_params = self._get_datadog_llm_obs_params()
            kwargs.update(dict_datadog_llm_obs_params)
            CustomBatchLogger.__init__(self, **kwargs, flush_lock=self.flush_lock)
        except Exception as e:
            verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}")
            raise e

    def _configure_dd_agent(self, dd_agent_host: str):
        """
        Configure the Datadog logger to send traces to the Agent.
        """
        # When using the Agent, LLM Observability Intake does NOT require the API Key
        # Reference: https://docs.datadoghq.com/llm_observability/setup/sdk/#agent-setup

        # Use specific port for LLM Obs (Trace Agent) to avoid conflict with Logs Agent (10518)
        agent_port = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126")
        self.DD_SITE = "localhost"  # Not used for URL construction in agent mode
        self.intake_url = (
            f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans"
        )
        verbose_logger.debug(f"DataDogLLMObs: Using DD Agent at {self.intake_url}")

    def _configure_dd_direct_api(self):
        """
        Configure the Datadog logger to send traces directly to the Datadog API.
        """
        if not self.DD_API_KEY:
            raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'")

        self.DD_SITE = os.getenv("DD_SITE")
        if not self.DD_SITE:
            raise Exception(
                "DD_SITE is not set, set 'DD_SITE=<>', example site = `us5.datadoghq.com`"
            )

        self.intake_url = (
            f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans"
        )

    def _get_datadog_llm_obs_params(self) -> Dict:
        """
        Get the datadog_llm_observability_params from litellm.datadog_llm_observability_params

        These are params specific to initializing the DataDogLLMObsLogger e.g. turn_off_message_logging
        """
        dict_datadog_llm_obs_params: Dict = {}
        if litellm.datadog_llm_observability_params is not None:
            if isinstance(
                litellm.datadog_llm_observability_params, DatadogLLMObsInitParams
            ):
                dict_datadog_llm_obs_params = (
                    litellm.datadog_llm_observability_params.model_dump()
                )
            elif isinstance(litellm.datadog_llm_observability_params, Dict):
                # only allow params that are of DatadogLLMObsInitParams
                dict_datadog_llm_obs_params = DatadogLLMObsInitParams(
                    **litellm.datadog_llm_observability_params
                ).model_dump()
        return dict_datadog_llm_obs_params

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        try:
            verbose_logger.debug(
                f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}"
            )
            payload = self.create_llm_obs_payload(kwargs, start_time, end_time)
            verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}")
            self.log_queue.append(payload)

            if len(self.log_queue) >= self.batch_size:
                await self.async_send_batch()
        except Exception as e:
            verbose_logger.exception(
                f"DataDogLLMObs: Error logging success event - {str(e)}"
            )

    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
        try:
            verbose_logger.debug(
                f"DataDogLLMObs: Logging failure event for model {kwargs.get('model', 'unknown')}"
            )
            payload = self.create_llm_obs_payload(kwargs, start_time, end_time)
            verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}")
            self.log_queue.append(payload)

            if len(self.log_queue) >= self.batch_size:
                await self.async_send_batch()
        except Exception as e:
            verbose_logger.exception(
                f"DataDogLLMObs: Error logging failure event - {str(e)}"
            )

    async def async_send_batch(self):
        try:
            if not self.log_queue:
                return

            verbose_logger.debug(
                f"DataDogLLMObs: Flushing {len(self.log_queue)} events"
            )

            if self.is_mock_mode:
                verbose_logger.debug(
                    "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted"
                )

            # Prepare the payload
            payload = {
                "data": DDIntakePayload(
                    type="span",
                    attributes=DDSpanAttributes(
                        ml_app=get_datadog_service(),
                        tags=[get_datadog_tags()],
                        spans=self.log_queue,
                    ),
                ),
            }

            # serialize datetime objects - for budget reset time in spend metrics
            from litellm.litellm_core_utils.safe_json_dumps import safe_dumps

            try:
                verbose_logger.debug("payload %s", safe_dumps(payload))
            except Exception as debug_error:
                verbose_logger.debug(
                    "payload serialization failed: %s", str(debug_error)
                )

            json_payload = safe_dumps(payload)

            headers = {"Content-Type": "application/json"}
            if self.DD_API_KEY:
                headers["DD-API-KEY"] = self.DD_API_KEY

            response = await self.async_client.post(
                url=self.intake_url,
                content=json_payload,
                headers=headers,
            )

            if response.status_code != 202:
                raise Exception(
                    f"DataDogLLMObs: Unexpected response - status_code: {response.status_code}, text: {response.text}"
                )

            if self.is_mock_mode:
                verbose_logger.debug(
                    f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked"
                )
            else:
                verbose_logger.debug(
                    f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}"
                )
            self.log_queue.clear()
        except httpx.HTTPStatusError as e:
            verbose_logger.exception(
                f"DataDogLLMObs: Error sending batch - {e.response.text}"
            )
        except Exception as e:
            verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {str(e)}")

    def create_llm_obs_payload(
        self, kwargs: Dict, start_time: datetime, end_time: datetime
    ) -> LLMObsPayload:
        standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
            "standard_logging_object"
        )
        if standard_logging_payload is None:
            raise Exception("DataDogLLMObs: standard_logging_object is not set")

        messages = standard_logging_payload["messages"]
        messages = self._ensure_string_content(messages=messages)

        metadata = kwargs.get("litellm_params", {}).get("metadata", {})

        input_meta = InputMeta(
            messages=handle_any_messages_to_chat_completion_str_messages_conversion(
                messages
            )
        )
        output_meta = OutputMeta(
            messages=self._get_response_messages(
                standard_logging_payload=standard_logging_payload,
                call_type=standard_logging_payload.get("call_type"),
            )
        )

        error_info = self._assemble_error_info(standard_logging_payload)

        metadata_parent_id: Optional[str] = None
        if isinstance(metadata, dict):
            metadata_parent_id = metadata.get("parent_id")

        meta = Meta(
            kind=self._get_datadog_span_kind(
                standard_logging_payload.get("call_type"), metadata_parent_id
            ),
            input=input_meta,
            output=output_meta,
            metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload),
            error=error_info,
        )

        # Calculate metrics (you may need to adjust these based on available data)
        metrics = LLMMetrics(
            input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)),
            output_tokens=float(standard_logging_payload.get("completion_tokens", 0)),
            total_tokens=float(standard_logging_payload.get("total_tokens", 0)),
            total_cost=float(standard_logging_payload.get("response_cost", 0)),
            time_to_first_token=self._get_time_to_first_token_seconds(
                standard_logging_payload
            ),
        )

        payload: LLMObsPayload = LLMObsPayload(
            parent_id=metadata_parent_id if metadata_parent_id else "undefined",
            trace_id=standard_logging_payload.get("trace_id", str(uuid.uuid4())),
            span_id=metadata.get("span_id", str(uuid.uuid4())),
            name=metadata.get("name", "litellm_llm_call"),
            meta=meta,
            start_ns=int(start_time.timestamp() * 1e9),
            duration=int((end_time - start_time).total_seconds() * 1e9),
            metrics=metrics,
            status="error" if error_info else "ok",
            tags=[get_datadog_tags(standard_logging_object=standard_logging_payload)],
        )

        apm_trace_id = self._get_apm_trace_id()
        if apm_trace_id is not None:
            payload["apm_id"] = apm_trace_id

        return payload

    def _get_apm_trace_id(self) -> Optional[str]:
        """Retrieve the current APM trace ID if available."""
        try:
            current_span_fn = getattr(tracer, "current_span", None)
            if callable(current_span_fn):
                current_span = current_span_fn()
                if current_span is not None:
                    trace_id = getattr(current_span, "trace_id", None)
                    if trace_id is not None:
                        return str(trace_id)
        except Exception:
            pass
        return None

    def _assemble_error_info(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> Optional[DDLLMObsError]:
        """
        Assemble error information for failure cases according to DD LLM Obs API spec
        """
        # Handle error information for failure cases according to DD LLM Obs API spec
        error_info: Optional[DDLLMObsError] = None

        if standard_logging_payload.get("status") == "failure":
            # Try to get structured error information first
            error_information: Optional[
                StandardLoggingPayloadErrorInformation
            ] = standard_logging_payload.get("error_information")

            if error_information:
                error_info = DDLLMObsError(
                    message=error_information.get("error_message")
                    or standard_logging_payload.get("error_str")
                    or "Unknown error",
                    type=error_information.get("error_class"),
                    stack=error_information.get("traceback"),
                )
        return error_info

    def _get_time_to_first_token_seconds(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> float:
        """
        Get the time to first token in seconds

        CompletionStartTime - StartTime = Time to first token

        For non streaming calls, CompletionStartTime is time we get the response back
        """
        start_time: Optional[float] = standard_logging_payload.get("startTime")
        completion_start_time: Optional[float] = standard_logging_payload.get(
            "completionStartTime"
        )
        end_time: Optional[float] = standard_logging_payload.get("endTime")

        if completion_start_time is not None and start_time is not None:
            return completion_start_time - start_time
        elif end_time is not None and start_time is not None:
            return end_time - start_time
        else:
            return 0.0

    def _get_response_messages(
        self, standard_logging_payload: StandardLoggingPayload, call_type: Optional[str]
    ) -> List[Any]:
        """
        Get the messages from the response object

        for now this handles logging /chat/completions responses
        """

        response_obj = standard_logging_payload.get("response")
        if response_obj is None:
            return []

        # edge case: handle response_obj is a string representation of a dict
        if isinstance(response_obj, str):
            try:
                import ast

                response_obj = ast.literal_eval(response_obj)
            except (ValueError, SyntaxError):
                try:
                    # fallback to json parsing
                    response_obj = json.loads(str(response_obj))
                except json.JSONDecodeError:
                    return []

        if call_type in [
            CallTypes.completion.value,
            CallTypes.acompletion.value,
            CallTypes.text_completion.value,
            CallTypes.atext_completion.value,
            CallTypes.generate_content.value,
            CallTypes.agenerate_content.value,
            CallTypes.generate_content_stream.value,
            CallTypes.agenerate_content_stream.value,
            CallTypes.anthropic_messages.value,
        ]:
            try:
                # Safely extract message from response_obj, handle failure cases
                if isinstance(response_obj, dict) and "choices" in response_obj:
                    choices = response_obj["choices"]
                    if choices and len(choices) > 0 and "message" in choices[0]:
                        return [choices[0]["message"]]
                return []
            except (KeyError, IndexError, TypeError):
                # In case of any error accessing the response structure, return empty list
                return []
        return []

    def _get_datadog_span_kind(
        self, call_type: Optional[str], parent_id: Optional[str] = None
    ) -> Literal["llm", "tool", "task", "embedding", "retrieval"]:
        """
        Map liteLLM call_type to appropriate DataDog LLM Observability span kind.

        Available DataDog span kinds: "llm", "tool", "task", "embedding", "retrieval"
        see: https://docs.datadoghq.com/ja/llm_observability/terms/
        """
        # Non llm/workflow/agent kinds cannot be root spans, so fallback to "llm" when parent metadata is missing
        if call_type is None or parent_id is None:
            return "llm"

        # Embedding operations
        if call_type in [CallTypes.embedding.value, CallTypes.aembedding.value]:
            return "embedding"

        # LLM completion operations
        if call_type in [
            CallTypes.completion.value,
            CallTypes.acompletion.value,
            CallTypes.text_completion.value,
            CallTypes.atext_completion.value,
            CallTypes.generate_content.value,
            CallTypes.agenerate_content.value,
            CallTypes.generate_content_stream.value,
            CallTypes.agenerate_content_stream.value,
            CallTypes.anthropic_messages.value,
            CallTypes.responses.value,
            CallTypes.aresponses.value,
        ]:
            return "llm"

        # Tool operations
        if call_type in [CallTypes.call_mcp_tool.value]:
            return "tool"

        # Retrieval operations
        if call_type in [
            CallTypes.get_assistants.value,
            CallTypes.aget_assistants.value,
            CallTypes.get_thread.value,
            CallTypes.aget_thread.value,
            CallTypes.get_messages.value,
            CallTypes.aget_messages.value,
            CallTypes.afile_retrieve.value,
            CallTypes.file_retrieve.value,
            CallTypes.afile_list.value,
            CallTypes.file_list.value,
            CallTypes.afile_content.value,
            CallTypes.file_content.value,
            CallTypes.retrieve_batch.value,
            CallTypes.aretrieve_batch.value,
            CallTypes.retrieve_fine_tuning_job.value,
            CallTypes.aretrieve_fine_tuning_job.value,
            CallTypes.alist_input_items.value,
        ]:
            return "retrieval"

        # Task operations (batch, fine-tuning, file operations, etc.)
        if call_type in [
            CallTypes.create_batch.value,
            CallTypes.acreate_batch.value,
            CallTypes.create_fine_tuning_job.value,
            CallTypes.acreate_fine_tuning_job.value,
            CallTypes.cancel_fine_tuning_job.value,
            CallTypes.acancel_fine_tuning_job.value,
            CallTypes.list_fine_tuning_jobs.value,
            CallTypes.alist_fine_tuning_jobs.value,
            CallTypes.create_assistants.value,
            CallTypes.acreate_assistants.value,
            CallTypes.delete_assistant.value,
            CallTypes.adelete_assistant.value,
            CallTypes.create_thread.value,
            CallTypes.acreate_thread.value,
            CallTypes.add_message.value,
            CallTypes.a_add_message.value,
            CallTypes.run_thread.value,
            CallTypes.arun_thread.value,
            CallTypes.run_thread_stream.value,
            CallTypes.arun_thread_stream.value,
            CallTypes.file_delete.value,
            CallTypes.afile_delete.value,
            CallTypes.create_file.value,
            CallTypes.acreate_file.value,
            CallTypes.image_generation.value,
            CallTypes.aimage_generation.value,
            CallTypes.image_edit.value,
            CallTypes.aimage_edit.value,
            CallTypes.moderation.value,
            CallTypes.amoderation.value,
            CallTypes.transcription.value,
            CallTypes.atranscription.value,
            CallTypes.speech.value,
            CallTypes.aspeech.value,
            CallTypes.rerank.value,
            CallTypes.arerank.value,
        ]:
            return "task"

        # Default fallback for unknown or passthrough operations
        return "llm"

    def _ensure_string_content(
        self, messages: Optional[Union[str, List[Any], Dict[Any, Any]]]
    ) -> List[Any]:
        if messages is None:
            return []
        if isinstance(messages, str):
            return [messages]
        elif isinstance(messages, list):
            return [message for message in messages]
        elif isinstance(messages, dict):
            return [str(messages.get("content", ""))]
        return []

    def _get_dd_llm_obs_payload_metadata(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> Dict[str, Any]:
        """
        Fields to track in DD LLM Observability metadata from litellm standard logging payload
        """
        _metadata: Dict[str, Any] = {
            "model_name": standard_logging_payload.get("model", "unknown"),
            "model_provider": standard_logging_payload.get(
                "custom_llm_provider", "unknown"
            ),
            "id": standard_logging_payload.get("id", "unknown"),
            "trace_id": standard_logging_payload.get("trace_id", "unknown"),
            "cache_hit": standard_logging_payload.get("cache_hit", "unknown"),
            "cache_key": standard_logging_payload.get("cache_key", "unknown"),
            "saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0),
            "guardrail_information": standard_logging_payload.get(
                "guardrail_information", None
            ),
            "is_streamed_request": self._get_stream_value_from_payload(
                standard_logging_payload
            ),
        }

        #########################################################
        # Add latency metrics to metadata
        #########################################################
        latency_metrics = self._get_latency_metrics(standard_logging_payload)
        _metadata.update({"latency_metrics": dict(latency_metrics)})

        #########################################################
        # Add spend metrics to metadata
        #########################################################
        spend_metrics = self._get_spend_metrics(standard_logging_payload)
        _metadata.update({"spend_metrics": dict(spend_metrics)})

        ## extract tool calls and add to metadata
        tool_call_metadata = self._extract_tool_call_metadata(standard_logging_payload)
        _metadata.update(tool_call_metadata)

        _standard_logging_metadata: dict = (
            dict(standard_logging_payload.get("metadata", {})) or {}
        )
        _metadata.update(_standard_logging_metadata)
        return _metadata

    def _get_latency_metrics(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> DDLLMObsLatencyMetrics:
        """
        Get the latency metrics from the standard logging payload
        """
        latency_metrics: DDLLMObsLatencyMetrics = DDLLMObsLatencyMetrics()
        # Add latency metrics to metadata
        # Time to first token (convert from seconds to milliseconds for consistency)
        time_to_first_token_seconds = self._get_time_to_first_token_seconds(
            standard_logging_payload
        )
        if time_to_first_token_seconds > 0:
            latency_metrics["time_to_first_token_ms"] = (
                time_to_first_token_seconds * 1000
            )

        # LiteLLM overhead time
        hidden_params = standard_logging_payload.get("hidden_params", {})
        litellm_overhead_ms = hidden_params.get("litellm_overhead_time_ms")
        if litellm_overhead_ms is not None:
            latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms

        # Guardrail overhead latency
        guardrail_info: Optional[
            list[StandardLoggingGuardrailInformation]
        ] = standard_logging_payload.get("guardrail_information")
        if guardrail_info is not None:
            total_duration = 0.0
            for info in guardrail_info:
                _guardrail_duration_seconds: Optional[float] = info.get("duration")
                if _guardrail_duration_seconds is not None:
                    total_duration += float(_guardrail_duration_seconds)

            if total_duration > 0:
                # Convert from seconds to milliseconds for consistency
                latency_metrics["guardrail_overhead_time_ms"] = total_duration * 1000

        return latency_metrics

    def _get_stream_value_from_payload(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> bool:
        """
        Extract the stream value from standard logging payload.

        The stream field in StandardLoggingPayload is only set to True for completed streaming responses.
        For non-streaming requests, it's None. The original stream parameter is in model_parameters.

        Returns:
            bool: True if this was a streaming request, False otherwise
        """
        # Check top-level stream field first (only True for completed streaming)
        stream_value = standard_logging_payload.get("stream")
        if stream_value is True:
            return True

        # Fallback to model_parameters.stream for original request parameters
        model_params = standard_logging_payload.get("model_parameters", {})
        if isinstance(model_params, dict):
            stream_value = model_params.get("stream")
            if stream_value is True:
                return True

        # Default to False for non-streaming requests
        return False

    def _get_spend_metrics(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> DDLLMObsSpendMetrics:
        """
        Get the spend metrics from the standard logging payload
        """
        spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics()

        # send response cost
        spend_metrics["response_cost"] = standard_logging_payload.get(
            "response_cost", 0.0
        )

        # Get budget information from metadata
        metadata = standard_logging_payload.get("metadata", {})

        # API key max budget
        user_api_key_max_budget = metadata.get("user_api_key_max_budget")
        if user_api_key_max_budget is not None:
            spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)

        # API key spend
        user_api_key_spend = metadata.get("user_api_key_spend")
        if user_api_key_spend is not None:
            try:
                spend_metrics["user_api_key_spend"] = float(user_api_key_spend)
            except (ValueError, TypeError):
                verbose_logger.debug(
                    f"Invalid user_api_key_spend value: {user_api_key_spend}"
                )

        # API key budget reset datetime
        user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at")
        if user_api_key_budget_reset_at is not None:
            try:
                from datetime import datetime, timezone

                budget_reset_at = None
                if isinstance(user_api_key_budget_reset_at, str):
                    # Handle ISO format strings that might have 'Z' suffix
                    iso_string = user_api_key_budget_reset_at.replace("Z", "+00:00")
                    budget_reset_at = datetime.fromisoformat(iso_string)
                elif isinstance(user_api_key_budget_reset_at, datetime):
                    budget_reset_at = user_api_key_budget_reset_at

                if budget_reset_at is not None:
                    # Preserve timezone info if already present
                    if budget_reset_at.tzinfo is None:
                        budget_reset_at = budget_reset_at.replace(tzinfo=timezone.utc)

                    # Convert to ISO string format for JSON serialization
                    # This prevents circular reference issues and ensures proper timezone representation
                    iso_string = budget_reset_at.isoformat()
                    spend_metrics["user_api_key_budget_reset_at"] = iso_string

                    # Debug logging to verify the conversion
                    verbose_logger.debug(
                        f"Converted budget_reset_at to ISO format: {iso_string}"
                    )
            except Exception as e:
                verbose_logger.debug(f"Error processing budget reset datetime: {e}")
                verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}")

        return spend_metrics

    def _process_input_messages_preserving_tool_calls(
        self, messages: List[Any]
    ) -> List[Dict[str, Any]]:
        """
        Process input messages while preserving tool_calls and tool message types.

        This bypasses the lossy string conversion when tool calls are present,
        allowing complex nested tool_calls objects to be preserved for Datadog.
        """
        processed = []
        for msg in messages:
            if isinstance(msg, dict):
                # Preserve messages with tool_calls or tool role as-is
                if "tool_calls" in msg or msg.get("role") == "tool":
                    processed.append(msg)
                else:
                    # For regular messages, still apply string conversion
                    converted = (
                        handle_any_messages_to_chat_completion_str_messages_conversion(
                            [msg]
                        )
                    )
                    processed.extend(converted)
            else:
                # For non-dict messages, apply string conversion
                converted = (
                    handle_any_messages_to_chat_completion_str_messages_conversion(
                        [msg]
                    )
                )
                processed.extend(converted)
        return processed

    @staticmethod
    def _tool_calls_kv_pair(tool_calls: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Extract tool call information into key-value pairs for Datadog metadata.

        Similar to OpenTelemetry's implementation but adapted for Datadog's format.
        """
        kv_pairs: Dict[str, Any] = {}
        for idx, tool_call in enumerate(tool_calls):
            try:
                # Extract tool call ID
                tool_id = tool_call.get("id")
                if tool_id:
                    kv_pairs[f"tool_calls.{idx}.id"] = tool_id

                # Extract tool call type
                tool_type = tool_call.get("type")
                if tool_type:
                    kv_pairs[f"tool_calls.{idx}.type"] = tool_type

                # Extract function information
                function = tool_call.get("function")
                if function:
                    function_name = function.get("name")
                    if function_name:
                        kv_pairs[f"tool_calls.{idx}.function.name"] = function_name

                    function_arguments = function.get("arguments")
                    if function_arguments:
                        # Store arguments as JSON string for Datadog
                        if isinstance(function_arguments, str):
                            kv_pairs[
                                f"tool_calls.{idx}.function.arguments"
                            ] = function_arguments
                        else:
                            import json

                            kv_pairs[
                                f"tool_calls.{idx}.function.arguments"
                            ] = json.dumps(function_arguments)
            except (KeyError, TypeError, ValueError) as e:
                verbose_logger.debug(
                    f"DataDogLLMObs: Error processing tool call {idx}: {str(e)}"
                )
                continue

        return kv_pairs

    def _extract_tool_call_metadata(
        self, standard_logging_payload: StandardLoggingPayload
    ) -> Dict[str, Any]:
        """
        Extract tool call information from both input messages and response for Datadog metadata.
        """
        tool_call_metadata: Dict[str, Any] = {}

        try:
            # Extract tool calls from input messages
            messages = standard_logging_payload.get("messages", [])
            if messages and isinstance(messages, list):
                for message in messages:
                    if isinstance(message, dict) and "tool_calls" in message:
                        tool_calls = message.get("tool_calls")
                        if tool_calls:
                            input_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
                            # Prefix with "input_" to distinguish from response tool calls
                            for key, value in input_tool_calls_kv.items():
                                tool_call_metadata[f"input_{key}"] = value

            # Extract tool calls from response
            response_obj = standard_logging_payload.get("response")
            if response_obj and isinstance(response_obj, dict):
                choices = response_obj.get("choices", [])
                for choice in choices:
                    if isinstance(choice, dict):
                        message = choice.get("message")
                        if message and isinstance(message, dict):
                            tool_calls = message.get("tool_calls")
                            if tool_calls:
                                response_tool_calls_kv = self._tool_calls_kv_pair(
                                    tool_calls
                                )
                                # Prefix with "output_" to distinguish from input tool calls
                                for key, value in response_tool_calls_kv.items():
                                    tool_call_metadata[f"output_{key}"] = value

        except Exception as e:
            verbose_logger.debug(
                f"DataDogLLMObs: Error extracting tool call metadata: {str(e)}"
            )

        return tool_call_metadata
