"""
OpenAI Chat Completions Message Handler for Unified Guardrails

This module provides a class-based handler for OpenAI-format chat completions.
The class methods can be overridden for custom behavior.

Pattern Overview:
-----------------
1. Extract text content from messages/responses (both string and list formats)
2. Create async tasks to apply guardrails to each text segment
3. Track mappings to know where each response belongs
4. Apply guardrail responses back to the original structure

This pattern can be replicated for other message formats (e.g., Anthropic).
"""

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.utils import (
    Choices,
    GenericGuardrailAPIInputs,
    ModelResponse,
    ModelResponseStream,
    StreamingChoices,
)

if TYPE_CHECKING:
    from litellm.integrations.custom_guardrail import CustomGuardrail


class OpenAIChatCompletionsHandler(BaseTranslation):
    """
    Handler for processing OpenAI chat completions messages with guardrails.

    This class provides methods to:
    1. Process input messages (pre-call hook)
    2. Process output responses (post-call hook)

    Methods can be overridden to customize behavior for different message formats.
    """

    async def process_input_messages(
        self,
        data: dict,
        guardrail_to_apply: "CustomGuardrail",
        litellm_logging_obj: Optional[Any] = None,
    ) -> Any:
        """
        Process input messages by applying guardrails to text content.
        """
        messages = data.get("messages")
        if messages is None:
            return data

        texts_to_check: List[str] = []
        images_to_check: List[str] = []
        tool_calls_to_check: List[ChatCompletionToolParam] = []
        text_task_mappings: List[Tuple[int, Optional[int]]] = []
        tool_call_task_mappings: List[Tuple[int, int]] = []
        # text_task_mappings: Track (message_index, content_index) for each text
        # content_index is None for string content, int for list content
        # tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call

        # Step 1: Extract all text content, images, and tool calls
        for msg_idx, message in enumerate(messages):
            self._extract_inputs(
                message=message,
                msg_idx=msg_idx,
                texts_to_check=texts_to_check,
                images_to_check=images_to_check,
                tool_calls_to_check=tool_calls_to_check,
                text_task_mappings=text_task_mappings,
                tool_call_task_mappings=tool_call_task_mappings,
            )

        # Step 2: Apply guardrail to all texts and tool calls in batch
        if texts_to_check or tool_calls_to_check:
            inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
            if images_to_check:
                inputs["images"] = images_to_check
            if tool_calls_to_check:
                inputs["tool_calls"] = tool_calls_to_check  # type: ignore
            if messages:
                inputs[
                    "structured_messages"
                ] = messages  # pass the openai /chat/completions messages to the guardrail, as-is
            # Pass tools (function definitions) to the guardrail
            tools = data.get("tools")
            if tools:
                inputs["tools"] = tools
            # Include model information if available
            model = data.get("model")
            if model:
                inputs["model"] = model

            guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
                inputs=inputs,
                request_data=data,
                input_type="request",
                logging_obj=litellm_logging_obj,
            )

            guardrailed_texts = guardrailed_inputs.get("texts", [])
            guardrailed_tool_calls = guardrailed_inputs.get("tool_calls", [])
            guardrailed_tools = guardrailed_inputs.get("tools")
            if guardrailed_tools is not None:
                data["tools"] = guardrailed_tools

            # Step 3: Map guardrail responses back to original message structure
            if guardrailed_texts and texts_to_check:
                await self._apply_guardrail_responses_to_input_texts(
                    messages=messages,
                    responses=guardrailed_texts,
                    task_mappings=text_task_mappings,
                )

            # Step 4: Apply guardrailed tool calls back to messages
            if guardrailed_tool_calls:
                # Note: The guardrail may modify tool_calls_to_check in place
                # or we may need to handle returned tool calls differently
                await self._apply_guardrail_responses_to_input_tool_calls(
                    messages=messages,
                    tool_calls=guardrailed_tool_calls,  # type: ignore
                    task_mappings=tool_call_task_mappings,
                )

        verbose_proxy_logger.debug(
            "OpenAI Chat Completions: Processed input messages: %s", messages
        )

        return data

    def extract_request_tool_names(self, data: dict) -> List[str]:
        """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
        names: List[str] = []
        for tool in data.get("tools") or []:
            if isinstance(tool, dict) and tool.get("type") == "function":
                fn = tool.get("function")
                if isinstance(fn, dict) and fn.get("name"):
                    names.append(str(fn["name"]))
        for fn in data.get("functions") or []:
            if isinstance(fn, dict) and fn.get("name"):
                names.append(str(fn["name"]))
        return names

    def _extract_inputs(
        self,
        message: Dict[str, Any],
        msg_idx: int,
        texts_to_check: List[str],
        images_to_check: List[str],
        tool_calls_to_check: List[ChatCompletionToolParam],
        text_task_mappings: List[Tuple[int, Optional[int]]],
        tool_call_task_mappings: List[Tuple[int, int]],
    ) -> None:
        """
        Extract text content, images, and tool calls from a message.

        Override this method to customize text/image/tool call extraction logic.
        """
        content = message.get("content", None)
        if content is not None:
            if isinstance(content, str):
                # Simple string content
                texts_to_check.append(content)
                text_task_mappings.append((msg_idx, None))

            elif isinstance(content, list):
                # List content (e.g., multimodal with text and images)
                for content_idx, content_item in enumerate(content):
                    # Extract text
                    text_str = content_item.get("text", None)
                    if text_str is not None:
                        texts_to_check.append(text_str)
                        text_task_mappings.append((msg_idx, int(content_idx)))

                    # Extract images (image_url)
                    if content_item.get("type") == "image_url":
                        image_url = content_item.get("image_url", {})
                        if isinstance(image_url, dict):
                            url = image_url.get("url")
                            if url:
                                images_to_check.append(url)
                        elif isinstance(image_url, str):
                            images_to_check.append(image_url)

        # Extract tool calls (typically in assistant messages)
        tool_calls = message.get("tool_calls", None)
        if tool_calls is not None and isinstance(tool_calls, list):
            for tool_call_idx, tool_call in enumerate(tool_calls):
                if isinstance(tool_call, dict):
                    # Add the full tool call object to the list
                    tool_calls_to_check.append(cast(ChatCompletionToolParam, tool_call))
                    tool_call_task_mappings.append((msg_idx, int(tool_call_idx)))

    async def _apply_guardrail_responses_to_input_texts(
        self,
        messages: List[Dict[str, Any]],
        responses: List[str],
        task_mappings: List[Tuple[int, Optional[int]]],
    ) -> None:
        """
        Apply guardrail responses back to input message text content.

        Override this method to customize how text responses are applied.
        """
        for task_idx, guardrail_response in enumerate(responses):
            mapping = task_mappings[task_idx]
            msg_idx = cast(int, mapping[0])
            content_idx_optional = cast(Optional[int], mapping[1])

            # Handle content
            content = messages[msg_idx].get("content", None)
            if content is None:
                continue

            if isinstance(content, str) and content_idx_optional is None:
                # Replace string content with guardrail response
                messages[msg_idx]["content"] = guardrail_response

            elif isinstance(content, list) and content_idx_optional is not None:
                # Replace specific text item in list content
                messages[msg_idx]["content"][content_idx_optional][
                    "text"
                ] = guardrail_response

    async def _apply_guardrail_responses_to_input_tool_calls(
        self,
        messages: List[Dict[str, Any]],
        tool_calls: List[Dict[str, Any]],
        task_mappings: List[Tuple[int, int]],
    ) -> None:
        """
        Apply guardrailed tool calls back to input messages.

        The guardrail may have modified the tool_calls list in place,
        so we apply the modified tool calls back to the original messages.

        Override this method to customize how tool call responses are applied.
        """
        for task_idx, (msg_idx, tool_call_idx) in enumerate(task_mappings):
            if task_idx < len(tool_calls):
                guardrailed_tool_call = tool_calls[task_idx]
                message_tool_calls = messages[msg_idx].get("tool_calls", None)
                if message_tool_calls is not None and isinstance(
                    message_tool_calls, list
                ):
                    if tool_call_idx < len(message_tool_calls):
                        # Replace the tool call with the guardrailed version
                        message_tool_calls[tool_call_idx] = guardrailed_tool_call

    async def process_output_response(
        self,
        response: "ModelResponse",
        guardrail_to_apply: "CustomGuardrail",
        litellm_logging_obj: Optional[Any] = None,
        user_api_key_dict: Optional[Any] = None,
    ) -> Any:
        """
        Process output response by applying guardrails to text content.

        Args:
            response: LiteLLM ModelResponse object
            guardrail_to_apply: The guardrail instance to apply
            litellm_logging_obj: Optional logging object
            user_api_key_dict: User API key metadata to pass to guardrails

        Returns:
            Modified response with guardrail applied to content

        Response Format Support:
            - String content: choice.message.content = "text here"
            - List content: choice.message.content = [{"type": "text", "text": "text here"}, ...]
        """

        # Step 0: Check if response has any text content to process
        if not self._has_text_content(response):
            verbose_proxy_logger.warning(
                "OpenAI Chat Completions: No text content in response, skipping guardrail"
            )
            return response

        texts_to_check: List[str] = []
        images_to_check: List[str] = []
        tool_calls_to_check: List[Dict[str, Any]] = []
        text_task_mappings: List[Tuple[int, Optional[int]]] = []
        tool_call_task_mappings: List[Tuple[int, int]] = []
        # text_task_mappings: Track (choice_index, content_index) for each text
        # content_index is None for string content, int for list content
        # tool_call_task_mappings: Track (choice_index, tool_call_index) for each tool call

        # Step 1: Extract all text content, images, and tool calls from response choices
        for choice_idx, choice in enumerate(response.choices):
            self._extract_output_text_images_and_tool_calls(
                choice=choice,
                choice_idx=choice_idx,
                texts_to_check=texts_to_check,
                images_to_check=images_to_check,
                tool_calls_to_check=tool_calls_to_check,
                text_task_mappings=text_task_mappings,
                tool_call_task_mappings=tool_call_task_mappings,
            )

        # Step 2: Apply guardrail to all texts and tool calls in batch
        if texts_to_check or tool_calls_to_check:
            # Create a request_data dict with response info and user API key metadata
            request_data: dict = {"response": response}

            # Add user API key metadata with prefixed keys
            user_metadata = self.transform_user_api_key_dict_to_metadata(
                user_api_key_dict
            )
            if user_metadata:
                request_data["litellm_metadata"] = user_metadata

            inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
            if images_to_check:
                inputs["images"] = images_to_check
            if tool_calls_to_check:
                inputs["tool_calls"] = tool_calls_to_check  # type: ignore
            # Include model information from the response if available
            if hasattr(response, "model") and response.model:
                inputs["model"] = response.model

            guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
                inputs=inputs,
                request_data=request_data,
                input_type="response",
                logging_obj=litellm_logging_obj,
            )

            guardrailed_texts = guardrailed_inputs.get("texts", [])

            # Step 3: Map guardrail responses back to original response structure
            if guardrailed_texts and texts_to_check:
                await self._apply_guardrail_responses_to_output_texts(
                    response=response,
                    responses=guardrailed_texts,
                    task_mappings=text_task_mappings,
                )

            # Step 4: Apply guardrailed tool calls back to response
            if tool_calls_to_check:
                await self._apply_guardrail_responses_to_output_tool_calls(
                    response=response,
                    tool_calls=tool_calls_to_check,
                    task_mappings=tool_call_task_mappings,
                )

        verbose_proxy_logger.debug(
            "OpenAI Chat Completions: Processed output response: %s", response
        )

        return response

    async def process_output_streaming_response(
        self,
        responses_so_far: List["ModelResponseStream"],
        guardrail_to_apply: "CustomGuardrail",
        litellm_logging_obj: Optional[Any] = None,
        user_api_key_dict: Optional[Any] = None,
    ) -> List["ModelResponseStream"]:
        """
        Process output streaming responses by applying guardrails to text content.

        Args:
            responses_so_far: List of LiteLLM ModelResponseStream objects
            guardrail_to_apply: The guardrail instance to apply
            litellm_logging_obj: Optional logging object
            user_api_key_dict: User API key metadata to pass to guardrails

        Returns:
            Modified list of responses with guardrail applied to content

        Response Format Support:
            - String content: choice.message.content = "text here"
            - List content: choice.message.content = [{"type": "text", "text": "text here"}, ...]
        """
        # check if the stream has ended
        has_stream_ended = False
        for chunk in responses_so_far:
            if chunk.choices and chunk.choices[0].finish_reason is not None:
                has_stream_ended = True
                break

        if has_stream_ended:
            # convert to model response
            model_response = cast(
                ModelResponse,
                stream_chunk_builder(
                    chunks=responses_so_far, logging_obj=litellm_logging_obj
                ),
            )
            # run process_output_response
            await self.process_output_response(
                response=model_response,
                guardrail_to_apply=guardrail_to_apply,
                litellm_logging_obj=litellm_logging_obj,
                user_api_key_dict=user_api_key_dict,
            )

            return responses_so_far

        # Step 0: Check if any response has text content to process
        has_any_text_content = False
        for response in responses_so_far:
            if self._has_text_content(response):
                has_any_text_content = True
                break

        if not has_any_text_content:
            verbose_proxy_logger.warning(
                "OpenAI Chat Completions: No text content in streaming responses, skipping guardrail"
            )
            return responses_so_far

        # Step 1: Combine all streaming chunks into complete text per choice
        # For streaming, we need to concatenate all delta.content across all chunks
        # Key: (choice_idx, content_idx), Value: combined text
        combined_texts = self._combine_streaming_texts(responses_so_far)

        # Step 2: Create lists for guardrail processing
        texts_to_check: List[str] = []
        images_to_check: List[str] = []
        task_mappings: List[Tuple[int, Optional[int]]] = []
        # Track (choice_index, content_index) for each combined text

        for (map_choice_idx, map_content_idx), combined_text in combined_texts.items():
            texts_to_check.append(combined_text)
            task_mappings.append((map_choice_idx, map_content_idx))

        # Step 3: Apply guardrail to all combined texts in batch
        if texts_to_check:
            # Create a request_data dict with response info and user API key metadata
            request_data: dict = {"responses": responses_so_far}

            # Add user API key metadata with prefixed keys
            user_metadata = self.transform_user_api_key_dict_to_metadata(
                user_api_key_dict
            )
            if user_metadata:
                request_data["litellm_metadata"] = user_metadata

            inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
            if images_to_check:
                inputs["images"] = images_to_check
            # Include model information from the first response if available
            if (
                responses_so_far
                and hasattr(responses_so_far[0], "model")
                and responses_so_far[0].model
            ):
                inputs["model"] = responses_so_far[0].model
            guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
                inputs=inputs,
                request_data=request_data,
                input_type="response",
                logging_obj=litellm_logging_obj,
            )

            guardrailed_texts = guardrailed_inputs.get("texts", [])

            # Step 4: Apply guardrailed text back to all streaming chunks
            # For each choice, replace the combined text across all chunks
            await self._apply_guardrail_responses_to_output_streaming(
                responses=responses_so_far,
                guardrailed_texts=guardrailed_texts,
                task_mappings=task_mappings,
            )

        verbose_proxy_logger.debug(
            "OpenAI Chat Completions: Processed output streaming responses: %s",
            responses_so_far,
        )

        return responses_so_far

    def _combine_streaming_texts(
        self, responses_so_far: List["ModelResponseStream"]
    ) -> Dict[Tuple[int, Optional[int]], str]:
        """
        Combine all streaming chunks into complete text per choice.

        For streaming, we need to concatenate all delta.content across all chunks.

        Args:
            responses_so_far: List of LiteLLM ModelResponseStream objects

        Returns:
            Dict mapping (choice_idx, content_idx) to combined text string
        """
        combined_texts: Dict[Tuple[int, Optional[int]], str] = {}

        for response_idx, response in enumerate(responses_so_far):
            for choice_idx, choice in enumerate(response.choices):
                if isinstance(choice, litellm.StreamingChoices):
                    content = choice.delta.content
                elif isinstance(choice, litellm.Choices):
                    content = choice.message.content
                else:
                    continue

                if content is None:
                    continue

                if isinstance(content, str):
                    # String content - accumulate for this choice
                    str_key: Tuple[int, Optional[int]] = (choice_idx, None)
                    if str_key not in combined_texts:
                        combined_texts[str_key] = ""
                    combined_texts[str_key] += content

                elif isinstance(content, list):
                    # List content - accumulate for each content item
                    for content_idx, content_item in enumerate(content):
                        text_str = content_item.get("text")
                        if text_str:
                            list_key: Tuple[int, Optional[int]] = (
                                choice_idx,
                                content_idx,
                            )
                            if list_key not in combined_texts:
                                combined_texts[list_key] = ""
                            combined_texts[list_key] += text_str

        return combined_texts

    def _has_text_content(
        self, response: Union["ModelResponse", "ModelResponseStream"]
    ) -> bool:
        """
        Check if response has any text content or tool calls to process.

        Override this method to customize text content detection.
        """
        from litellm.types.utils import ModelResponse, ModelResponseStream

        if isinstance(response, ModelResponse):
            for choice in response.choices:
                if isinstance(choice, litellm.Choices):
                    # Check for text content
                    if choice.message.content and isinstance(
                        choice.message.content, str
                    ):
                        return True
                    # Check for tool calls
                    if choice.message.tool_calls and isinstance(
                        choice.message.tool_calls, list
                    ):
                        if len(choice.message.tool_calls) > 0:
                            return True
        elif isinstance(response, ModelResponseStream):
            for streaming_choice in response.choices:
                if isinstance(streaming_choice, litellm.StreamingChoices):
                    # Check for text content
                    if streaming_choice.delta.content and isinstance(
                        streaming_choice.delta.content, str
                    ):
                        return True
                    # Check for tool calls
                    if streaming_choice.delta.tool_calls and isinstance(
                        streaming_choice.delta.tool_calls, list
                    ):
                        if len(streaming_choice.delta.tool_calls) > 0:
                            return True
        return False

    def _extract_output_text_images_and_tool_calls(
        self,
        choice: Union[Choices, StreamingChoices],
        choice_idx: int,
        texts_to_check: List[str],
        images_to_check: List[str],
        tool_calls_to_check: List[Dict[str, Any]],
        text_task_mappings: List[Tuple[int, Optional[int]]],
        tool_call_task_mappings: List[Tuple[int, int]],
    ) -> None:
        """
        Extract text content, images, and tool calls from a response choice.

        Override this method to customize text/image/tool call extraction logic.
        """
        verbose_proxy_logger.debug(
            "OpenAI Chat Completions: Processing choice: %s", choice
        )

        # Determine content source and tool calls based on choice type
        content = None
        tool_calls: Optional[List[Any]] = None
        if isinstance(choice, litellm.Choices):
            content = choice.message.content
            tool_calls = choice.message.tool_calls
        elif isinstance(choice, litellm.StreamingChoices):
            content = choice.delta.content
            tool_calls = choice.delta.tool_calls
        else:
            # Unknown choice type, skip processing
            return

        # Process content if it exists
        if content and isinstance(content, str):
            # Simple string content
            texts_to_check.append(content)
            text_task_mappings.append((choice_idx, None))

        elif content and isinstance(content, list):
            # List content (e.g., multimodal response)
            for content_idx, content_item in enumerate(content):
                # Extract text
                content_text = content_item.get("text")
                if content_text:
                    texts_to_check.append(content_text)
                    text_task_mappings.append((choice_idx, int(content_idx)))

                # Extract images
                if content_item.get("type") == "image_url":
                    image_url = content_item.get("image_url", {})
                    if isinstance(image_url, dict):
                        url = image_url.get("url")
                        if url:
                            images_to_check.append(url)

        # Process tool calls if they exist
        if tool_calls is not None and isinstance(tool_calls, list):
            for tool_call_idx, tool_call in enumerate(tool_calls):
                # Convert tool call to dict format for guardrail processing
                tool_call_dict = self._convert_tool_call_to_dict(tool_call)
                if tool_call_dict:
                    tool_calls_to_check.append(tool_call_dict)
                    tool_call_task_mappings.append((choice_idx, int(tool_call_idx)))

    def _convert_tool_call_to_dict(
        self, tool_call: Union[Dict[str, Any], Any]
    ) -> Optional[Dict[str, Any]]:
        """
        Convert a tool call object to dictionary format.

        Tool calls can be either dict or object depending on the type.
        """
        if isinstance(tool_call, dict):
            return tool_call
        elif hasattr(tool_call, "id") and hasattr(tool_call, "function"):
            # Convert object to dict
            function = tool_call.function
            function_dict = {}
            if hasattr(function, "name"):
                function_dict["name"] = function.name
            if hasattr(function, "arguments"):
                function_dict["arguments"] = function.arguments

            tool_call_dict = {
                "id": tool_call.id if hasattr(tool_call, "id") else None,
                "type": tool_call.type if hasattr(tool_call, "type") else "function",
                "function": function_dict,
            }
            return tool_call_dict
        return None

    async def _apply_guardrail_responses_to_output_texts(
        self,
        response: "ModelResponse",
        responses: List[str],
        task_mappings: List[Tuple[int, Optional[int]]],
    ) -> None:
        """
        Apply guardrail text responses back to output response.

        Override this method to customize how text responses are applied.
        """
        for task_idx, guardrail_response in enumerate(responses):
            mapping = task_mappings[task_idx]
            choice_idx = cast(int, mapping[0])
            content_idx_optional = cast(Optional[int], mapping[1])

            choice = cast(Choices, response.choices[choice_idx])

            # Handle content
            content = choice.message.content
            if content is None:
                continue

            if isinstance(content, str) and content_idx_optional is None:
                # Replace string content with guardrail response
                choice.message.content = guardrail_response

            elif isinstance(content, list) and content_idx_optional is not None:
                # Replace specific text item in list content
                choice.message.content[content_idx_optional]["text"] = guardrail_response  # type: ignore

    async def _apply_guardrail_responses_to_output_tool_calls(
        self,
        response: "ModelResponse",
        tool_calls: List[Dict[str, Any]],
        task_mappings: List[Tuple[int, int]],
    ) -> None:
        """
        Apply guardrailed tool calls back to output response.

        The guardrail may have modified the tool_calls list in place,
        so we apply the modified tool calls back to the original response.

        Override this method to customize how tool call responses are applied.
        """
        for task_idx, (choice_idx, tool_call_idx) in enumerate(task_mappings):
            if task_idx < len(tool_calls):
                guardrailed_tool_call = tool_calls[task_idx]
                choice = cast(Choices, response.choices[choice_idx])
                choice_tool_calls = choice.message.tool_calls

                if choice_tool_calls is not None and isinstance(
                    choice_tool_calls, list
                ):
                    if tool_call_idx < len(choice_tool_calls):
                        # Update the tool call with guardrailed version
                        existing_tool_call = choice_tool_calls[tool_call_idx]
                        # Update object attributes (output responses always have typed objects)
                        if "function" in guardrailed_tool_call:
                            func_dict = guardrailed_tool_call["function"]
                            if "arguments" in func_dict:
                                existing_tool_call.function.arguments = func_dict[
                                    "arguments"
                                ]
                            if "name" in func_dict:
                                existing_tool_call.function.name = func_dict["name"]

    async def _apply_guardrail_responses_to_output_streaming(
        self,
        responses: List["ModelResponseStream"],
        guardrailed_texts: List[str],
        task_mappings: List[Tuple[int, Optional[int]]],
    ) -> None:
        """
        Apply guardrail responses back to output streaming responses.

        For streaming responses, the guardrailed text (which is the combined text from all chunks)
        is placed in the first chunk, and subsequent chunks are cleared.

        Args:
            responses: List of ModelResponseStream objects to modify
            guardrailed_texts: List of guardrailed text responses (combined from all chunks)
            task_mappings: List of tuples (choice_idx, content_idx)

        Override this method to customize how responses are applied to streaming responses.
        """
        # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx)
        guardrail_map: Dict[Tuple[int, Optional[int]], str] = {}
        for task_idx, guardrail_response in enumerate(guardrailed_texts):
            mapping = task_mappings[task_idx]
            choice_idx = cast(int, mapping[0])
            content_idx_optional = cast(Optional[int], mapping[1])
            guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response

        # Track which choices we've already set the guardrailed text for
        # Key: (choice_idx, content_idx), Value: boolean (True if already set)
        already_set: Dict[Tuple[int, Optional[int]], bool] = {}

        # Iterate through all responses and update content
        for response_idx, response in enumerate(responses):
            for choice_idx_in_response, choice in enumerate(response.choices):
                if isinstance(choice, litellm.StreamingChoices):
                    content = choice.delta.content
                elif isinstance(choice, litellm.Choices):
                    content = choice.message.content
                else:
                    continue

                if content is None:
                    continue

                if isinstance(content, str):
                    # String content
                    str_key: Tuple[int, Optional[int]] = (choice_idx_in_response, None)
                    if str_key in guardrail_map:
                        if str_key not in already_set:
                            # First chunk - set the complete guardrailed text
                            if isinstance(choice, litellm.StreamingChoices):
                                choice.delta.content = guardrail_map[str_key]
                            elif isinstance(choice, litellm.Choices):
                                choice.message.content = guardrail_map[str_key]
                            already_set[str_key] = True
                        else:
                            # Subsequent chunks - clear the content
                            if isinstance(choice, litellm.StreamingChoices):
                                choice.delta.content = ""
                            elif isinstance(choice, litellm.Choices):
                                choice.message.content = ""

                elif isinstance(content, list):
                    # List content - handle each content item
                    for content_idx, content_item in enumerate(content):
                        if "text" in content_item:
                            list_key: Tuple[int, Optional[int]] = (
                                choice_idx_in_response,
                                content_idx,
                            )
                            if list_key in guardrail_map:
                                if list_key not in already_set:
                                    # First chunk - set the complete guardrailed text
                                    content_item["text"] = guardrail_map[list_key]
                                    already_set[list_key] = True
                                else:
                                    # Subsequent chunks - clear the text
                                    content_item["text"] = ""
