"""
Unit tests for BaseResponsesAPIStreamingIterator

Tests core functionality including:
1. Processing chunks and handling ResponseCompletedEvent 
2. Ensuring _update_responses_api_response_id_with_model_id is called for final chunk
3. Verifying ID update is NOT called for non-final chunks (delta events)
4. Edge case handling for invalid JSON, empty chunks, and [DONE] markers

These tests ensure the streaming iterator correctly processes response chunks 
and applies model ID updates only to completed responses, as required for proper
response tracking and logging.
"""

import json
from datetime import datetime
from typing import Any, Dict, Optional
from unittest.mock import Mock, patch

import pytest


from litellm.constants import STREAM_SSE_DONE_STRING
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
    ResponseCompletedEvent,
    ResponseFailedEvent,
    ResponseIncompleteEvent,
    ResponsesAPIResponse,
    ResponsesAPIStreamEvents,
    OutputTextDeltaEvent,
)


class TestBaseResponsesAPIStreamingIterator:
    """Test cases for BaseResponsesAPIStreamingIterator"""

    @pytest.mark.asyncio
    async def test_responses_streaming_iterator_parses_u2028_in_sse_json(self):
        """
        U+2028 inside JSON must not split the SSE event. httpx aiter_lines uses
        str.splitlines() and drops response.completed; OpenAI SSEDecoder does not.
        """
        from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator

        u2028 = "\u2028"
        payload = json.dumps(
            {
                "type": "response.completed",
                "response": {"instructions": f"eligible{u2028}promo"},
            }
        )
        sse_bytes = f"data: {payload}\n\n".encode("utf-8")

        async def mock_aiter_bytes():
            yield sse_bytes

        mock_response = Mock()
        mock_response.headers = {}
        mock_response.aiter_bytes = mock_aiter_bytes

        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
        mock_responses_api_response.id = "resp_u2028"
        mock_completed_event = Mock(spec=ResponseCompletedEvent)
        mock_completed_event.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
        mock_completed_event.response = mock_responses_api_response
        mock_config.transform_streaming_response.return_value = mock_completed_event

        iterator = ResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        chunks = []
        with (
            patch("asyncio.create_task"),
            patch("litellm.responses.streaming_iterator.executor"),
        ):
            async for chunk in iterator:
                chunks.append(chunk)

        assert len(chunks) == 1
        assert chunks[0].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
        assert iterator.completed_response is not None

    def test_process_chunk_with_response_completed_event(self):
        """
        Test that _process_chunk correctly processes a ResponseCompletedEvent
        and calls _update_responses_api_response_id_with_model_id for the final chunk.
        """
        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        # Create a mock ResponsesAPIResponse for the completed event
        mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
        mock_responses_api_response.id = "original_response_id"

        # Create a mock ResponseCompletedEvent
        mock_completed_event = Mock(spec=ResponseCompletedEvent)
        mock_completed_event.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
        mock_completed_event.response = mock_responses_api_response

        # Set up the mock transform method to return our completed event
        mock_config.transform_streaming_response.return_value = mock_completed_event

        # Mock the _update_responses_api_response_id_with_model_id method
        updated_response = Mock(spec=ResponsesAPIResponse)
        updated_response.id = "updated_response_id"

        # Create the iterator instance
        iterator = BaseResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        # Prepare test chunk data
        test_chunk_data = {
            "type": "response.completed",
            "response": {
                "id": "original_response_id",
                "output": [{"type": "message", "content": [{"text": "Hello World"}]}],
            },
        }

        with patch.object(
            ResponsesAPIRequestUtils,
            "_update_responses_api_response_id_with_model_id",
            return_value=updated_response,
        ) as mock_update_id:
            # Process the chunk
            result = iterator._process_chunk(json.dumps(test_chunk_data))

            # Assertions
            assert result is not None
            assert result.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED

            # Verify that _update_responses_api_response_id_with_model_id was called
            mock_update_id.assert_called_once_with(
                responses_api_response=mock_responses_api_response,
                litellm_metadata={"model_info": {"id": "model_123"}},
                custom_llm_provider="openai",
            )

            # Verify the completed response was stored
            assert iterator.completed_response == result

            # Verify the response was updated on the event
            assert result.response == updated_response

    def test_process_chunk_with_delta_event_no_id_update(self):
        """
        Test that _process_chunk correctly processes a delta event
        and does NOT call _update_responses_api_response_id_with_model_id.
        """
        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        # Create a mock OutputTextDeltaEvent (not a completed event)
        mock_delta_event = Mock(spec=OutputTextDeltaEvent)
        mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA
        mock_delta_event.delta = "Hello"
        # Delta events don't have a response attribute
        (
            delattr(mock_delta_event, "response")
            if hasattr(mock_delta_event, "response")
            else None
        )

        # Set up the mock transform method to return our delta event
        mock_config.transform_streaming_response.return_value = mock_delta_event

        # Create the iterator instance
        iterator = BaseResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        # Prepare test chunk data for a delta event
        test_chunk_data = {
            "type": "response.output_text.delta",
            "delta": "Hello",
            "item_id": "item_123",
            "output_index": 0,
            "content_index": 0,
        }

        with patch.object(
            ResponsesAPIRequestUtils, "_update_responses_api_response_id_with_model_id"
        ) as mock_update_id:
            # Process the chunk
            result = iterator._process_chunk(json.dumps(test_chunk_data))

            # Assertions
            assert result is not None
            assert result.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA

            # Verify that _update_responses_api_response_id_with_model_id was NOT called
            mock_update_id.assert_not_called()

            # Verify no completed response was stored (since this is not a completed event)
            assert iterator.completed_response is None

    def test_process_chunk_handles_invalid_json(self):
        """
        Test that _process_chunk gracefully handles invalid JSON.
        """
        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        # Create the iterator instance
        iterator = BaseResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
        )

        # Test with invalid JSON
        result = iterator._process_chunk("invalid json {")

        # Should return None for invalid JSON
        assert result is None
        assert iterator.completed_response is None

    def test_process_chunk_handles_done_marker(self):
        """
        Test that _process_chunk correctly handles the [DONE] marker.
        """
        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        # Create the iterator instance
        iterator = BaseResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
        )

        # Test with [DONE] marker
        result = iterator._process_chunk(STREAM_SSE_DONE_STRING)

        # Should return None and set finished flag
        assert result is None
        assert iterator.finished is True

    def test_process_chunk_handles_empty_chunk(self):
        """
        Test that _process_chunk correctly handles empty or None chunks.
        """
        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        # Create the iterator instance
        iterator = BaseResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
        )

        # Test with empty chunk
        result = iterator._process_chunk("")
        assert result is None

        # Test with None chunk
        result = iterator._process_chunk(None)
        assert result is None

    def test_handle_logging_completed_response_with_unpickleable_objects(self):
        """
        Test that _handle_logging_completed_response handles responses containing
        objects that cannot be pickled (like Pydantic ValidatorIterator).

        This test verifies the fix for issue #17192 where streaming with tool_choice
        containing allowed_tools would fail with:
        "cannot pickle 'pydantic_core._pydantic_core.ValidatorIterator' object"

        The fix uses model_dump + model_validate instead of copy.deepcopy.
        """
        import asyncio
        from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator

        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}
        mock_response.aiter_bytes = Mock()
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_logging_obj.async_success_handler = Mock()
        mock_logging_obj.success_handler = Mock()
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        # Create the iterator instance
        iterator = ResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        # Create a ResponseCompletedEvent with tool_choice that has model_dump
        mock_completed_response = Mock()
        mock_completed_response.model_dump.return_value = {
            "type": "response.completed",
            "response": {
                "id": "resp_123",
                "output": [{"type": "function_call", "name": "search_web"}],
                "tool_choice": {"type": "function", "name": "search_web"},
            },
        }
        # model_validate should return a new mock (the copy)
        type(mock_completed_response).model_validate = Mock(return_value=Mock())

        iterator.completed_response = mock_completed_response

        # This should NOT raise an exception
        # Previously it would fail with: TypeError: cannot pickle 'ValidatorIterator'
        # Mock asyncio.create_task and executor.submit since we're not in async context
        with (
            patch("asyncio.create_task") as mock_create_task,
            patch("litellm.responses.streaming_iterator.executor") as mock_executor,
        ):
            try:
                iterator._handle_logging_completed_response()
            except TypeError as e:
                if "pickle" in str(e):
                    pytest.fail(
                        f"_handle_logging_completed_response failed with pickle error: {e}"
                    )
                raise

    @pytest.mark.asyncio
    async def test_stop_async_iteration_not_logged_as_failure(self):
        """
        Test that StopAsyncIteration is NOT logged as a failure.

        This test verifies that when streaming completes normally with StopAsyncIteration,
        the _handle_failure method is NOT called, preventing false error logs in Langfuse
        and other logging integrations.

        """
        from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator

        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}

        async def mock_aiter_bytes():
            yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n'

        mock_response.aiter_bytes = mock_aiter_bytes

        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_logging_obj.async_failure_handler = Mock()
        mock_logging_obj.failure_handler = Mock()

        mock_config = Mock(spec=BaseResponsesAPIConfig)
        mock_delta_event = Mock()
        mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA
        mock_delta_event.delta = "test"
        mock_config.transform_streaming_response.return_value = mock_delta_event

        # Create the iterator instance
        iterator = ResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        # Consume the iterator until StopAsyncIteration
        chunks_received = []
        try:
            async for chunk in iterator:
                chunks_received.append(chunk)
        except StopAsyncIteration:
            pass  # This is expected

        # Verify we got the chunk
        assert len(chunks_received) == 1

        # CRITICAL: Verify that failure handlers were NOT called
        # StopAsyncIteration is a normal end of stream, not a failure
        mock_logging_obj.async_failure_handler.assert_not_called()
        mock_logging_obj.failure_handler.assert_not_called()

    def test_stop_iteration_not_logged_as_failure(self):
        """
        Test that StopIteration is NOT logged as a failure in sync iterator.

        This test verifies that when streaming completes normally with StopIteration,
        the _handle_failure method is NOT called, preventing false error logs in Langfuse
        and other logging integrations.

        Regression test for: https://github.com/BerriAI/litellm/issues/XXXXX
        """
        from litellm.responses.streaming_iterator import (
            SyncResponsesAPIStreamingIterator,
        )

        # Mock dependencies
        mock_response = Mock()
        mock_response.headers = {}

        def mock_iter_bytes():
            yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n'

        mock_response.iter_bytes = mock_iter_bytes

        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_logging_obj.async_failure_handler = Mock()
        mock_logging_obj.failure_handler = Mock()

        mock_config = Mock(spec=BaseResponsesAPIConfig)
        mock_delta_event = Mock()
        mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA
        mock_delta_event.delta = "test"
        mock_config.transform_streaming_response.return_value = mock_delta_event

        # Create the iterator instance
        iterator = SyncResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        # Consume the iterator until StopIteration
        chunks_received = []
        try:
            for chunk in iterator:
                chunks_received.append(chunk)
        except StopIteration:
            pass  # This is expected

        # Verify we got the chunk
        assert len(chunks_received) == 1

        # CRITICAL: Verify that failure handlers were NOT called
        # StopIteration is a normal end of stream, not a failure
        mock_logging_obj.async_failure_handler.assert_not_called()
        mock_logging_obj.failure_handler.assert_not_called()

    def test_process_chunk_response_failed_calls_failure_handler(self):
        """
        Test that a RESPONSE_FAILED event routes to failure handlers,
        not success handlers. Failed responses represent genuine LLM-level
        errors and should be logged as failures.
        """
        from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator

        mock_response = Mock()
        mock_response.headers = {}
        mock_response.aiter_bytes = Mock()
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_logging_obj.async_failure_handler = Mock()
        mock_logging_obj.failure_handler = Mock()
        mock_logging_obj.async_success_handler = Mock()
        mock_logging_obj.success_handler = Mock()
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
        mock_responses_api_response.id = "resp_failed_123"
        mock_responses_api_response.error = {
            "type": "server_error",
            "message": "The model encountered an error",
        }
        mock_responses_api_response.usage = None

        mock_failed_event = Mock(spec=ResponseFailedEvent)
        mock_failed_event.type = ResponsesAPIStreamEvents.RESPONSE_FAILED
        mock_failed_event.response = mock_responses_api_response

        mock_config.transform_streaming_response.return_value = mock_failed_event

        iterator = ResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        test_chunk_data = {
            "type": "response.failed",
            "response": {
                "id": "resp_failed_123",
                "error": {
                    "type": "server_error",
                    "message": "The model encountered an error",
                },
            },
        }

        with (
            patch.object(
                ResponsesAPIRequestUtils,
                "_update_responses_api_response_id_with_model_id",
                return_value=mock_responses_api_response,
            ),
            patch(
                "litellm.responses.streaming_iterator.run_async_function"
            ) as mock_run_async,
            patch("litellm.responses.streaming_iterator.executor") as mock_executor,
        ):
            result = iterator._process_chunk(json.dumps(test_chunk_data))

            assert result is not None
            assert result.type == ResponsesAPIStreamEvents.RESPONSE_FAILED
            assert iterator.completed_response == result

            # Failure handler should have been called via _handle_failure
            mock_run_async.assert_called_once()
            call_kwargs = mock_run_async.call_args
            assert (
                call_kwargs[1]["async_function"]
                == mock_logging_obj.async_failure_handler
            )

            mock_executor.submit.assert_called_once()
            submit_args = mock_executor.submit.call_args
            assert submit_args[0][0] == mock_logging_obj.failure_handler

    def test_process_chunk_response_incomplete_calls_success_handler(self):
        """
        Test that a RESPONSE_INCOMPLETE event routes to success handlers.
        Incomplete responses (e.g. max_output_tokens reached) are still valid
        responses with usage data — analogous to finish_reason='length' in chat.
        """
        from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator

        mock_response = Mock()
        mock_response.headers = {}
        mock_response.aiter_bytes = Mock()
        mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
        mock_logging_obj.model_call_details = {"litellm_params": {}}
        mock_logging_obj.completion_start_time = None
        mock_logging_obj.async_failure_handler = Mock()
        mock_logging_obj.failure_handler = Mock()
        mock_logging_obj.async_success_handler = Mock()
        mock_logging_obj.success_handler = Mock()
        mock_config = Mock(spec=BaseResponsesAPIConfig)

        mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
        mock_responses_api_response.id = "resp_incomplete_123"
        mock_responses_api_response.incomplete_details = {"reason": "max_output_tokens"}
        mock_responses_api_response.usage = None

        mock_incomplete_event = Mock(spec=ResponseIncompleteEvent)
        mock_incomplete_event.type = ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE
        mock_incomplete_event.response = mock_responses_api_response

        mock_config.transform_streaming_response.return_value = mock_incomplete_event

        iterator = ResponsesAPIStreamingIterator(
            response=mock_response,
            model="gpt-5.5",
            responses_api_provider_config=mock_config,
            logging_obj=mock_logging_obj,
            litellm_metadata={"model_info": {"id": "model_123"}},
            custom_llm_provider="openai",
        )

        test_chunk_data = {
            "type": "response.incomplete",
            "response": {
                "id": "resp_incomplete_123",
                "incomplete_details": {"reason": "max_output_tokens"},
            },
        }

        with (
            patch.object(
                ResponsesAPIRequestUtils,
                "_update_responses_api_response_id_with_model_id",
                return_value=mock_responses_api_response,
            ),
            patch("asyncio.create_task") as mock_create_task,
            patch("litellm.responses.streaming_iterator.executor") as mock_executor,
        ):
            result = iterator._process_chunk(json.dumps(test_chunk_data))

            assert result is not None
            assert result.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE
            assert iterator.completed_response == result

            # Success handlers are dispatched as one async task (via _handle_logging_completed_response);
            # the sync handler must never be submitted to the executor concurrently (LIT-4210)
            mock_create_task.assert_called_once()
            mock_executor.submit.assert_not_called()

            # Failure handlers should NOT have been called
            mock_logging_obj.async_failure_handler.assert_not_called()
            mock_logging_obj.failure_handler.assert_not_called()
