import datetime
import os
import sys
import types
import unittest
from typing import Optional
from unittest.mock import MagicMock, patch

import pytest

import litellm
from litellm.integrations.langfuse import langfuse as langfuse_module
from litellm.integrations.langfuse.langfuse import LangFuseLogger

sys.path.insert(0, os.path.abspath("../.."))
from litellm.integrations.langfuse.langfuse import LangFuseLogger

# Import LangfuseUsageDetails directly from the module where it's defined
from litellm.types.integrations.langfuse import *


class TestLangfuseUsageDetails(unittest.TestCase):
    def setUp(self):
        # Save global Langfuse client counter to restore after test
        self._original_langfuse_clients_count = litellm.initialized_langfuse_clients

        # Set up environment variables for testing
        self.env_patcher = patch.dict(
            "os.environ",
            {
                "LANGFUSE_SECRET_KEY": "test-secret-key",
                "LANGFUSE_PUBLIC_KEY": "test-public-key",
                "LANGFUSE_HOST": "https://test.langfuse.com",
            },
        )
        self.env_patcher.start()

        # Create mock objects
        self.mock_langfuse_client = MagicMock()
        # Mock the client attribute to prevent errors during logger initialization
        self.mock_langfuse_client.client = MagicMock()
        self.mock_langfuse_trace = MagicMock()
        self.mock_langfuse_generation = MagicMock()
        self.mock_langfuse_generation.trace_id = "test-trace-id"

        # Mock span method for trace (used by log_provider_specific_information_as_span and _log_guardrail_information_as_span)
        self.mock_langfuse_span = MagicMock()
        self.mock_langfuse_span.end = MagicMock()
        self.mock_langfuse_trace.span.return_value = self.mock_langfuse_span

        # Setup the trace and generation chain
        self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation
        self.last_trace_kwargs = {}

        def _trace_side_effect(*args, **kwargs):
            self.last_trace_kwargs = kwargs
            return self.mock_langfuse_trace

        self.mock_langfuse_client.trace.side_effect = _trace_side_effect

        # Mock the langfuse module that's imported locally in methods
        self.langfuse_module_patcher = patch.dict(
            "sys.modules", {"langfuse": MagicMock()}
        )
        self.mock_langfuse_module = self.langfuse_module_patcher.start()

        # Create a mock for the langfuse module with version
        self.mock_langfuse = MagicMock()
        self.mock_langfuse.version = MagicMock()
        self.mock_langfuse.version.__version__ = (
            "3.0.0"  # Set a version that supports all features
        )

        # Mock the Langfuse class
        self.mock_langfuse_class = MagicMock()
        self.mock_langfuse_class.return_value = self.mock_langfuse_client

        # Set up the sys.modules['langfuse'] mock
        sys.modules["langfuse"] = self.mock_langfuse
        sys.modules["langfuse"].Langfuse = self.mock_langfuse_class

        # Create a fresh logger instance for each test
        self.logger = LangFuseLogger()

        # Explicitly set the Langfuse client to our mock
        self.logger.Langfuse = self.mock_langfuse_client
        # Ensure langfuse_sdk_version is set correctly for _supports_* methods
        self.logger.langfuse_sdk_version = "3.0.0"

        # Add the log_event_on_langfuse method to the instance
        def log_event_on_langfuse(
            self,
            kwargs,
            response_obj,
            start_time=None,
            end_time=None,
            user_id=None,
            level="DEFAULT",
            status_message=None,
        ):
            # This implementation calls _log_langfuse_v2 directly
            return self._log_langfuse_v2(
                user_id=user_id,
                metadata=kwargs.get("litellm_params", {}).get("metadata", {}),
                litellm_params=kwargs.get("litellm_params", {}),
                output=None,
                start_time=start_time,
                end_time=end_time,
                kwargs=kwargs,
                optional_params=kwargs.get("optional_params", {}),
                input=None,
                response_obj=response_obj,
                level=level,
                litellm_call_id=kwargs.get("litellm_call_id", None),
            )

        # Bind the method to the instance
        self.logger.log_event_on_langfuse = types.MethodType(
            log_event_on_langfuse, self.logger
        )

        # Make sure _is_langfuse_v2 returns True
        def mock_is_langfuse_v2(self):
            return True

        self.logger._is_langfuse_v2 = types.MethodType(mock_is_langfuse_v2, self.logger)

    def tearDown(self):
        # Clean up logger instance to prevent state leakage
        if hasattr(self, "logger"):
            # Reset logger's Langfuse client to break any references
            self.logger.Langfuse = None
            # Delete logger instance to ensure complete cleanup
            del self.logger

        # Restore global Langfuse client counter to prevent cross-test pollution
        litellm.initialized_langfuse_clients = self._original_langfuse_clients_count

        self.env_patcher.stop()
        self.langfuse_module_patcher.stop()  # patch.dict automatically restores sys.modules

    def test_langfuse_usage_details_type(self):
        """Test that LangfuseUsageDetails TypedDict is properly defined with the correct fields"""
        # Create an instance of LangfuseUsageDetails
        usage_details: LangfuseUsageDetails = {
            "input": 10,
            "output": 20,
            "total": 30,
            "cache_creation_input_tokens": 5,
            "cache_read_input_tokens": 3,
        }

        # Verify all fields are present
        self.assertEqual(usage_details["input"], 10)
        self.assertEqual(usage_details["output"], 20)
        self.assertEqual(usage_details["total"], 30)
        self.assertEqual(usage_details["cache_creation_input_tokens"], 5)
        self.assertEqual(usage_details["cache_read_input_tokens"], 3)

        # Test with all fields (all fields are required in TypedDict by default)
        minimal_usage_details: LangfuseUsageDetails = {
            "input": 10,
            "output": 20,
            "total": 30,
            "cache_creation_input_tokens": 0,
            "cache_read_input_tokens": 0,
        }

        self.assertEqual(minimal_usage_details["input"], 10)
        self.assertEqual(minimal_usage_details["output"], 20)
        self.assertEqual(minimal_usage_details["total"], 30)

    def test_log_langfuse_v2_usage_details(self):
        """Test that usage_details in _log_langfuse_v2 is correctly typed and assigned"""
        # Create a mock response object with usage information
        response_obj = MagicMock()
        response_obj.usage = MagicMock()
        response_obj.usage.prompt_tokens = 15
        response_obj.usage.completion_tokens = 25

        # Add the cache token attributes using get method
        def mock_get(key, default=None):
            if key == "cache_creation_input_tokens":
                return 7
            elif key == "cache_read_input_tokens":
                return 4
            return default

        response_obj.usage.get = mock_get

        # Create kwargs for the log_event method
        kwargs = {
            "model": "gpt-4",
            "messages": [{"role": "user", "content": "Hello"}],
            "litellm_params": {"metadata": {}},
        }

        # Create start and end times
        start_time = datetime.datetime.now()
        end_time = start_time + datetime.timedelta(seconds=1)

        # Call the log_event method
        with patch.object(self.logger, "_log_langfuse_v2") as mock_log_langfuse_v2:
            self.logger.log_event_on_langfuse(
                kwargs=kwargs,
                response_obj=response_obj,
                start_time=start_time,
                end_time=end_time,
            )

            # Check if _log_langfuse_v2 was called
            mock_log_langfuse_v2.assert_called_once()

            # Get the arguments passed to _log_langfuse_v2
            call_args = mock_log_langfuse_v2.call_args[1]

            # Verify response_obj was passed correctly
            self.assertEqual(call_args["response_obj"], response_obj)

    def test_langfuse_usage_details_optional_fields(self):
        """Test that LangfuseUsageDetails fields are properly defined as Optional"""
        # Create an instance with None values for optional fields
        usage_details: LangfuseUsageDetails = {
            "input": 10,
            "output": 20,
            "total": 30,
            "cache_creation_input_tokens": None,
            "cache_read_input_tokens": None,
        }

        # Verify fields can be None
        self.assertEqual(usage_details["input"], 10)
        self.assertEqual(usage_details["output"], 20)
        self.assertEqual(usage_details["total"], 30)
        self.assertIsNone(usage_details["cache_creation_input_tokens"])
        self.assertIsNone(usage_details["cache_read_input_tokens"])

    def test_langfuse_usage_details_structure(self):
        """Test that LangfuseUsageDetails has the correct structure as defined in the commit"""
        # This test directly verifies the structure of the TypedDict
        # without relying on the LangFuseLogger class

        # Create a dictionary that matches the LangfuseUsageDetails structure
        usage_details = {
            "input": 15,
            "output": 25,
            "total": 40,
            "cache_creation_input_tokens": 7,
            "cache_read_input_tokens": 4,
        }

        # Verify the structure matches what we expect
        self.assertIn("input", usage_details)
        self.assertIn("output", usage_details)
        self.assertIn("total", usage_details)
        self.assertIn("cache_creation_input_tokens", usage_details)
        self.assertIn("cache_read_input_tokens", usage_details)

        # Verify the values
        self.assertEqual(usage_details["input"], 15)
        self.assertEqual(usage_details["output"], 25)
        self.assertEqual(usage_details["total"], 40)
        self.assertEqual(usage_details["cache_creation_input_tokens"], 7)
        self.assertEqual(usage_details["cache_read_input_tokens"], 4)

    def test_log_langfuse_v2_handles_null_usage_values(self):
        """
        Test that _log_langfuse_v2 correctly handles None values in the usage object
        by converting them to 0, preventing validation errors.
        """
        # Reset the mock to ensure clean state; clear side_effect so return_value takes effect
        self.mock_langfuse_client.reset_mock(side_effect=True)
        self.mock_langfuse_trace.reset_mock(side_effect=True)
        self.mock_langfuse_generation.reset_mock(side_effect=True)

        # Re-setup the trace and generation chain with clean state
        self.mock_langfuse_generation.trace_id = "test-trace-id"
        mock_span = MagicMock()
        mock_span.end = MagicMock()
        self.mock_langfuse_trace.span.return_value = mock_span
        self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation

        # Ensure trace returns our mock
        self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace
        self.logger.Langfuse = self.mock_langfuse_client

        with (
            patch(
                "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
                side_effect=lambda generation_params, **kwargs: generation_params,
                create=True,
            ) as mock_add_prompt_params,
            patch.object(self.logger, "_supports_prompt", return_value=True),
        ):
            # Create a mock response object with usage information containing None values
            response_obj = MagicMock()
            response_obj.usage = MagicMock()
            response_obj.usage.prompt_tokens = None
            response_obj.usage.completion_tokens = None
            response_obj.usage.total_tokens = None

            # Mock the .get() method to return None for cache-related fields
            def mock_get(key, default=None):
                if key in ["cache_creation_input_tokens", "cache_read_input_tokens"]:
                    return None
                return default

            response_obj.usage.get = mock_get

            # Prepare standard kwargs for the call
            kwargs = {
                "model": "gpt-4-null-usage",
                "messages": [{"role": "user", "content": "Test"}],
                "litellm_params": {"metadata": {}},
                "optional_params": {},
                "litellm_call_id": "test-call-id-null-usage",
                "standard_logging_object": None,
                "response_cost": 0.0,
            }

            # Use fixed timestamps to avoid timing-related flakiness
            fixed_time = datetime.datetime(2024, 1, 1, 12, 0, 0)

            # Call the method under test
            try:
                self.logger._log_langfuse_v2(
                    user_id="test-user",
                    metadata={},
                    litellm_params=kwargs["litellm_params"],
                    output={"role": "assistant", "content": "Response"},
                    start_time=fixed_time,
                    end_time=fixed_time + datetime.timedelta(seconds=1),
                    kwargs=kwargs,
                    optional_params=kwargs["optional_params"],
                    input={"messages": kwargs["messages"]},
                    response_obj=response_obj,
                    level="DEFAULT",
                    litellm_call_id=kwargs["litellm_call_id"],
                )
            except Exception as e:
                self.fail(f"_log_langfuse_v2 raised an exception: {e}")

            # Verify that trace was called first
            self.mock_langfuse_client.trace.assert_called()

            #  Check the arguments passed to the mocked langfuse generation call
            self.mock_langfuse_trace.generation.assert_called_once()
            call_args, call_kwargs = self.mock_langfuse_trace.generation.call_args

            #  Inspect the usage and usage_details dictionaries
            usage_arg = call_kwargs.get("usage")
            usage_details_arg = call_kwargs.get("usage_details")

            self.assertIsNotNone(usage_arg)
            self.assertIsNotNone(usage_details_arg)

            # Verify that None values were converted to 0
            self.assertEqual(usage_arg["prompt_tokens"], 0)
            self.assertEqual(usage_arg["completion_tokens"], 0)

            self.assertEqual(usage_details_arg["input"], 0)
            self.assertEqual(usage_details_arg["output"], 0)
            self.assertEqual(usage_details_arg["total"], 0)
            self.assertEqual(usage_details_arg["cache_creation_input_tokens"], 0)
            self.assertEqual(usage_details_arg["cache_read_input_tokens"], 0)

            mock_add_prompt_params.assert_called_once()

    def _build_standard_logging_payload(self, trace_id: Optional[str] = None):
        payload = {
            "id": "payload-id",
            "call_type": "completion",
            "response_cost": 0.0,
            "status": "success",
            "total_tokens": 0,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "startTime": 0.0,
            "endTime": 0.0,
            "completionStartTime": 0.0,
            "model": "gpt-4",
            "model_id": "model-123",
            "model_group": "openai",
            "api_base": "https://api.openai.com",
            "metadata": {
                "user_api_key_end_user_id": None,
                "prompt_management_metadata": None,
                "session_id": None,
                "trace_name": None,
                "trace_version": None,
                "headers": None,
                "endpoint": None,
                "caching_groups": None,
                "previous_models": None,
            },
            "hidden_params": {},
            "request_tags": [],
            "messages": [],
            "response": {"id": "resp"},
            "model_parameters": {},
            "guardrail_information": None,
            "standard_built_in_tools_params": None,
        }
        if trace_id is not None:
            payload["trace_id"] = trace_id
        return payload

    def _build_langfuse_kwargs(self, standard_logging_payload):
        return {
            "standard_logging_object": standard_logging_payload,
            "model": standard_logging_payload["model"],
            "call_type": standard_logging_payload["call_type"],
            "cache_hit": False,
            "messages": [],
        }

    def test_log_langfuse_v2_uses_standard_trace_id_when_available(self):
        payload = self._build_standard_logging_payload(trace_id="std-trace-id")
        kwargs = self._build_langfuse_kwargs(payload)
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={},
                litellm_params={"metadata": {}},
                output=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input=None,
                response_obj=None,
                level="INFO",
                litellm_call_id="call-id-xyz",
            )

        assert self.last_trace_kwargs.get("id") == "std-trace-id"

    def test_log_langfuse_v2_defaults_to_call_id_without_standard_trace_id(self):
        payload = self._build_standard_logging_payload()
        kwargs = self._build_langfuse_kwargs(payload)
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={},
                litellm_params={"metadata": {}},
                output=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input=None,
                response_obj=None,
                level="INFO",
                litellm_call_id="call-id-xyz",
            )

        assert self.last_trace_kwargs.get("id") == "call-id-xyz"

    def test_log_langfuse_v2_uses_litellm_trace_id_fallback_over_call_id(self):
        """
        When standard_logging_object has no trace_id, but kwargs contains
        litellm_trace_id (the same ID the DB stores as Session ID), Langfuse
        should use litellm_trace_id — NOT litellm_call_id. This ensures the
        trace_id in Langfuse matches the Session ID shown in LiteLLM logs.
        """
        payload = self._build_standard_logging_payload()  # no trace_id
        kwargs = self._build_langfuse_kwargs(payload)
        kwargs["litellm_trace_id"] = "trace-id-from-kwargs"
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={},
                litellm_params={"metadata": {}},
                output=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input=None,
                response_obj=None,
                level="ERROR",
                litellm_call_id="call-id-xyz",
            )

        # litellm_trace_id should be preferred over litellm_call_id
        assert self.last_trace_kwargs.get("id") == "trace-id-from-kwargs"

    def test_log_langfuse_v2_uses_litellm_trace_id_when_standard_logging_object_none(
        self,
    ):
        """
        When standard_logging_object is None (failure case where
        get_standard_logging_object_payload threw), litellm_trace_id from kwargs
        should be used as the Langfuse trace_id. This matches the DB Session ID.
        """
        kwargs = {
            "standard_logging_object": None,
            "model": "gpt-4",
            "call_type": "completion",
            "cache_hit": False,
            "messages": [],
            "litellm_trace_id": "trace-id-failure",
        }
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={},
                litellm_params={"metadata": {}},
                output=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input=None,
                response_obj=None,
                level="ERROR",
                litellm_call_id="call-id-different",
            )

        # Must use litellm_trace_id, not litellm_call_id
        assert self.last_trace_kwargs.get("id") == "trace-id-failure"

    def test_log_langfuse_v2_session_id_passed_as_trace_session_id(self):
        """
        Test that metadata.session_id is correctly passed as trace_params["session_id"]
        for Langfuse session grouping, and does NOT override trace_id.
        Each LLM call should get its own unique trace_id while sharing the session_id.
        """
        payload = self._build_standard_logging_payload(trace_id="std-trace-123")
        kwargs = self._build_langfuse_kwargs(payload)
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={"session_id": "my-session-abc"},
                litellm_params={"metadata": {"session_id": "my-session-abc"}},
                output=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input=None,
                response_obj=None,
                level="INFO",
                litellm_call_id="call-id-456",
            )

        # session_id should be set for Langfuse session grouping
        assert self.last_trace_kwargs.get("session_id") == "my-session-abc"
        # trace_id should remain the standard trace_id, NOT the session_id
        assert self.last_trace_kwargs.get("id") == "std-trace-123"

    def test_log_langfuse_v2_session_id_preserved_for_error_level(self):
        """
        Test that session_id is correctly passed in trace_params even when
        the log level is ERROR (failure case). This verifies the fix for
        failed requests losing session_id mapping in Langfuse.
        """
        payload = self._build_standard_logging_payload(trace_id="std-trace-err")
        kwargs = self._build_langfuse_kwargs(payload)
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={"session_id": "error-session-xyz"},
                litellm_params={"metadata": {"session_id": "error-session-xyz"}},
                output="BadRequestError: model not found",
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input={"messages": [{"role": "user", "content": "test"}]},
                response_obj=None,
                level="ERROR",
                litellm_call_id="call-id-err-789",
            )

        # session_id must be preserved even for ERROR level logs
        assert self.last_trace_kwargs.get("session_id") == "error-session-xyz"
        # trace_id should be the standard trace_id, not the session_id
        assert self.last_trace_kwargs.get("id") == "std-trace-err"
        # status_message should be set for error traces
        assert self.last_trace_kwargs.get("status_message") is not None

    def test_log_langfuse_v2_explicit_trace_id_takes_priority_over_session_id(self):
        """
        Test that when both trace_id and session_id are provided in metadata,
        trace_id takes priority as the trace identifier.
        """
        payload = self._build_standard_logging_payload()
        kwargs = self._build_langfuse_kwargs(payload)
        self.last_trace_kwargs = {}

        with patch(
            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
            side_effect=lambda generation_params, **kwargs: generation_params,
            create=True,
        ):
            self.logger._log_langfuse_v2(
                user_id="user-1",
                metadata={
                    "session_id": "session-999",
                    "trace_id": "explicit-trace-id-777",
                },
                litellm_params={
                    "metadata": {
                        "session_id": "session-999",
                        "trace_id": "explicit-trace-id-777",
                    }
                },
                output=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
                kwargs=kwargs,
                optional_params={},
                input=None,
                response_obj=None,
                level="DEFAULT",
                litellm_call_id="call-id-aaa",
            )

        # Explicit trace_id must take priority
        assert self.last_trace_kwargs.get("id") == "explicit-trace-id-777"
        # session_id must still be set for session grouping
        assert self.last_trace_kwargs.get("session_id") == "session-999"


def test_failure_handler_langfuse_kwargs_excludes_original_response():
    """
    Test that the actual Logging.failure_handler() passes kwargs without
    'original_response' to the Langfuse logger. Exercises the real code path
    rather than simulating the filtering logic.
    """
    import litellm
    from litellm.litellm_core_utils.litellm_logging import Logging

    # Create a Logging instance
    logging_obj = Logging(
        model="gpt-4",
        messages=[{"role": "user", "content": "test"}],
        stream=False,
        call_type="completion",
        start_time=datetime.datetime.utcnow(),
        litellm_call_id="test-call-id-failure",
        function_id="test-function-id",
    )

    # Set up model_call_details with original_response (simulates a coroutine)
    mock_coroutine = MagicMock()
    logging_obj.model_call_details["original_response"] = mock_coroutine
    logging_obj.model_call_details["litellm_params"] = {
        "metadata": {"session_id": "test-session-failure"},
        "litellm_session_id": None,
    }
    logging_obj.model_call_details["optional_params"] = {}

    # Capture what gets passed to log_event_on_langfuse
    captured_kwargs = {}
    mock_langfuse_logger = MagicMock()

    def capture_log_event(**log_kwargs):
        captured_kwargs.update(log_kwargs)
        return {"trace_id": "mock-trace-id", "generation_id": "mock-gen-id"}

    mock_langfuse_logger.log_event_on_langfuse.side_effect = capture_log_event

    # Set "langfuse" as a failure callback so the failure_handler processes it
    original_failure_callback = litellm.failure_callback
    litellm.failure_callback = ["langfuse"]

    try:
        # Mock LangFuseHandler to return our capturing mock logger
        with patch(
            "litellm.litellm_core_utils.litellm_logging.LangFuseHandler"
        ) as mock_handler_class:
            mock_handler_class.get_langfuse_logger_for_request.return_value = (
                mock_langfuse_logger
            )

            # Call the actual failure_handler
            test_exception = Exception("TestError: model not found")
            logging_obj.failure_handler(
                exception=test_exception,
                traceback_exception="Traceback: test",
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
            )

        # Verify log_event_on_langfuse was actually called
        assert (
            mock_langfuse_logger.log_event_on_langfuse.called
        ), "log_event_on_langfuse was not called"

        # Verify original_response is NOT in the kwargs passed to Langfuse
        langfuse_kwargs = captured_kwargs.get("kwargs", {})
        assert (
            "original_response" not in langfuse_kwargs
        ), "original_response should be excluded from kwargs passed to Langfuse"

        # Verify session_id metadata is preserved in the kwargs
        langfuse_metadata = langfuse_kwargs.get("litellm_params", {}).get(
            "metadata", {}
        )
        assert (
            langfuse_metadata.get("session_id") == "test-session-failure"
        ), "session_id should be preserved in kwargs passed to Langfuse"

        # Verify level is ERROR
        assert captured_kwargs.get("level") == "ERROR"
    finally:
        litellm.failure_callback = original_failure_callback


@pytest.mark.asyncio
async def test_async_log_failure_event_logs_to_langfuse():
    """
    Test that LangfusePromptManagement.async_log_failure_event() calls
    log_event_on_langfuse with level=ERROR even when standard_logging_object
    is present. This is the code path the proxy uses for failed LLM calls.
    """
    from litellm.integrations.langfuse.langfuse_prompt_management import (
        LangfusePromptManagement,
    )

    mock_langfuse_module = MagicMock()
    mock_langfuse_module.version.__version__ = "3.0.0"

    with (
        patch.dict(
            "os.environ",
            {
                "LANGFUSE_SECRET_KEY": "test-secret",
                "LANGFUSE_PUBLIC_KEY": "test-public",
                "LANGFUSE_HOST": "https://test.langfuse.com",
            },
        ),
        patch.dict("sys.modules", {"langfuse": mock_langfuse_module}),
    ):
        prompt_mgmt = LangfusePromptManagement()

        # Mock the langfuse logger returned by get_langfuse_logger_for_request
        mock_logger = MagicMock()
        mock_logger.log_event_on_langfuse.return_value = {
            "trace_id": "mock-trace",
            "generation_id": "mock-gen",
        }

        with patch(
            "litellm.integrations.langfuse.langfuse_prompt_management.LangFuseHandler"
        ) as mock_handler:
            mock_handler.get_langfuse_logger_for_request.return_value = mock_logger

            kwargs = {
                "litellm_params": {
                    "metadata": {"session_id": "test-session-fail"},
                },
                "litellm_call_id": "call-fail-123",
                "user": "test-user",
                "exception": Exception("API error: model not found"),
                "standard_logging_object": {
                    "error_str": "API error: model not found",
                    "trace_id": "std-trace-fail",
                    "metadata": {},
                },
            }

            await prompt_mgmt.async_log_failure_event(
                kwargs=kwargs,
                response_obj=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
            )

            # Verify log_event_on_langfuse was called
            assert (
                mock_logger.log_event_on_langfuse.called
            ), "log_event_on_langfuse was not called for failure event"
            call_kwargs = mock_logger.log_event_on_langfuse.call_args[1]
            assert call_kwargs["level"] == "ERROR"
            assert call_kwargs["status_message"] == "API error: model not found"
            assert call_kwargs["response_obj"] is None


@pytest.mark.asyncio
async def test_async_log_failure_event_works_without_standard_logging_object():
    """
    Test that async_log_failure_event() still logs to Langfuse even when
    standard_logging_object is None (e.g. when get_standard_logging_object_payload
    threw an exception). This is the critical fix — before, it silently returned.
    """
    from litellm.integrations.langfuse.langfuse_prompt_management import (
        LangfusePromptManagement,
    )

    mock_langfuse_module = MagicMock()
    mock_langfuse_module.version.__version__ = "3.0.0"

    with (
        patch.dict(
            "os.environ",
            {
                "LANGFUSE_SECRET_KEY": "test-secret",
                "LANGFUSE_PUBLIC_KEY": "test-public",
                "LANGFUSE_HOST": "https://test.langfuse.com",
            },
        ),
        patch.dict("sys.modules", {"langfuse": mock_langfuse_module}),
    ):
        prompt_mgmt = LangfusePromptManagement()

        mock_logger = MagicMock()
        mock_logger.log_event_on_langfuse.return_value = {
            "trace_id": "mock-trace",
            "generation_id": "mock-gen",
        }

        with patch(
            "litellm.integrations.langfuse.langfuse_prompt_management.LangFuseHandler"
        ) as mock_handler:
            mock_handler.get_langfuse_logger_for_request.return_value = mock_logger

            kwargs = {
                "litellm_params": {
                    "metadata": {"session_id": "test-session-no-slo"},
                },
                "litellm_call_id": "call-no-slo-456",
                "user": "test-user",
                "exception": Exception("InternalServerError: something broke"),
                "standard_logging_object": None,  # This is the key — it's None
            }

            await prompt_mgmt.async_log_failure_event(
                kwargs=kwargs,
                response_obj=None,
                start_time=datetime.datetime.utcnow(),
                end_time=datetime.datetime.utcnow(),
            )

            # CRITICAL: log_event_on_langfuse MUST still be called
            assert mock_logger.log_event_on_langfuse.called, (
                "log_event_on_langfuse was NOT called when standard_logging_object "
                "is None — failure trace would be silently dropped"
            )
            call_kwargs = mock_logger.log_event_on_langfuse.call_args[1]
            assert call_kwargs["level"] == "ERROR"
            # Falls back to exception from kwargs
            assert "InternalServerError" in call_kwargs["status_message"]


def test_max_langfuse_clients_limit():
    """
    Test that the max langfuse clients limit is respected when initializing multiple clients
    """
    # Mock langfuse package to avoid triggering real import.
    # The real langfuse import fails on Python 3.14 due to pydantic v1 incompatibility,
    # and sys.modules["langfuse"] may be absent after other tests in the suite clean up.
    mock_langfuse = MagicMock()
    mock_langfuse.version.__version__ = "3.0.0"
    # Set max clients to 2 for testing
    original_initialized_langfuse_clients = litellm.initialized_langfuse_clients
    with (
        patch.dict("sys.modules", {"langfuse": mock_langfuse}),
        patch.object(langfuse_module, "MAX_LANGFUSE_INITIALIZED_CLIENTS", 2),
    ):
        # Reset the counter
        litellm.initialized_langfuse_clients = 0

        # First client should succeed
        logger1 = LangFuseLogger(
            langfuse_public_key="test_key_1",
            langfuse_secret="test_secret_1",
            langfuse_host="https://test1.langfuse.com",
        )
        assert litellm.initialized_langfuse_clients == 1

        # Second client should succeed
        logger2 = LangFuseLogger(
            langfuse_public_key="test_key_2",
            langfuse_secret="test_secret_2",
            langfuse_host="https://test2.langfuse.com",
        )
        assert litellm.initialized_langfuse_clients == 2

        # Third client should fail with exception
        with pytest.raises(Exception) as exc_info:
            logger3 = LangFuseLogger(
                langfuse_public_key="test_key_3",
                langfuse_secret="test_secret_3",
                langfuse_host="https://test3.langfuse.com",
            )

        # Verify the error message contains the expected text
        assert "Max langfuse clients reached" in str(exc_info.value)

        # Counter should still be 2 (third client failed to initialize)
        assert litellm.initialized_langfuse_clients == 2

    litellm.initialized_langfuse_clients = original_initialized_langfuse_clients
