import os
import sys
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch

import pytest

# newrelic is a proxy-runtime dependency (pyproject.toml) and is not installed
# in the CI Python environment. Mock it in sys.modules before importing the
# integration so that deferred `import newrelic.agent` calls inside NewRelicLogger
# methods resolve to these mocks rather than failing with ModuleNotFoundError.
_mock_newrelic = MagicMock()
_mock_newrelic_agent = MagicMock()
# Explicitly link so _mock_newrelic.agent IS _mock_newrelic_agent. Without this,
# the first getattr(_mock_newrelic, 'agent') auto-creates a different child mock,
# causing patch("newrelic.agent.xxx") to patch the wrong object.
_mock_newrelic.agent = _mock_newrelic_agent
sys.modules["newrelic"] = _mock_newrelic
sys.modules["newrelic.agent"] = _mock_newrelic_agent

import litellm
import litellm.integrations.newrelic.newrelic as nr_module
from litellm.integrations.newrelic.newrelic import NewRelicLogger


def _rendered_log_message(call):
    message = str(call.args[0])
    values = call.args[1:]
    return message % values if values else message


# The module may have been imported before sys.modules was patched (e.g. via
# litellm's own startup imports), leaving _newrelic_agent=None. Point it at
# the mock agent so all tests see a non-None agent.
nr_module._newrelic_agent = _mock_newrelic_agent


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

NR_ENV = {
    "NEW_RELIC_LICENSE_KEY": "test-license-key",
    "NEW_RELIC_APP_NAME": "test-app",
}


def make_logger(**kwargs) -> NewRelicLogger:
    """Instantiate NewRelicLogger with NR agent calls mocked out."""
    with patch.dict(os.environ, NR_ENV):
        return NewRelicLogger(**kwargs)


def make_kwargs(
    model="gpt-4",
    provider="openai",
    messages=None,
    optional_params=None,
    traceparent=None,
) -> dict:
    """Build a minimal kwargs dict representative of a litellm callback invocation."""
    headers = {}
    if traceparent:
        headers["traceparent"] = traceparent

    return {
        "model": model,
        "messages": messages or [{"role": "user", "content": "Hello"}],
        "optional_params": optional_params or {},
        "litellm_params": {
            "custom_llm_provider": provider,
            "metadata": {"headers": headers},
        },
        "start_time": 1_000_000.0,
        "end_time": 1_000_001.5,
        "llm_api_duration_ms": 1500.0,
    }


def make_response(
    model="gpt-4",
    response_id="chatcmpl-abc123",
    content="Hello there!",
    finish_reason="stop",
    prompt_tokens=10,
    completion_tokens=20,
):
    """Build a minimal ModelResponse-like dict."""
    return {
        "id": response_id,
        "model": model,
        "choices": [
            {
                "message": {"role": "assistant", "content": content},
                "finish_reason": finish_reason,
            }
        ],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }


def make_slo(**overrides):
    """Build a StandardLoggingPayload-like dict with sentinel values distinct from
    make_kwargs/make_response defaults, so tests can prove the SLO branch won."""
    base = {
        "trace_id": "slo-trace-abc",
        "custom_llm_provider": "slo-provider",
        "model": "slo-model",
        "prompt_tokens": 100,
        "completion_tokens": 200,
        "total_tokens": 300,
        "response_time": 1.5,  # seconds; converted to ms by _get_duration
        "model_parameters": {"temperature": 0.7, "max_tokens": 500},
        "startTime": 2_000_000.0,
        "endTime": 2_000_001.5,
        "messages": [{"role": "user", "content": "from-slo"}],
    }
    base.update(overrides)
    return base


# ---------------------------------------------------------------------------
# Init / configuration
# ---------------------------------------------------------------------------


class TestNewRelicLoggerInit:
    def test_disabled_when_license_key_missing(self):
        with patch("newrelic.agent.register_application"):
            with patch.dict(os.environ, {"NEW_RELIC_APP_NAME": "app"}, clear=True):
                logger = NewRelicLogger()
        assert logger.enabled is False

    def test_disabled_when_app_name_missing(self):
        with patch("newrelic.agent.register_application"):
            with patch.dict(os.environ, {"NEW_RELIC_LICENSE_KEY": "key"}, clear=True):
                logger = NewRelicLogger()
        assert logger.enabled is False

    def test_enabled_with_valid_env_vars(self):
        with patch("newrelic.agent.register_application"):
            with patch.dict(os.environ, NR_ENV):
                logger = NewRelicLogger()
        assert logger.enabled is True

    def test_disabled_on_import_error(self):
        with patch.object(
            _mock_newrelic_agent, "register_application", side_effect=ImportError
        ):
            with patch.dict(os.environ, NR_ENV):
                logger = NewRelicLogger()
        assert logger.enabled is False

    def test_disabled_on_agent_startup_error(self):
        with patch.object(
            _mock_newrelic_agent,
            "register_application",
            side_effect=RuntimeError("agent startup failed"),
        ):
            with patch.dict(os.environ, NR_ENV):
                logger = NewRelicLogger()
        assert logger.enabled is False

    def test_disabled_when_agent_package_missing(self):
        with patch.object(nr_module, "_newrelic_agent", None):
            with patch.dict(os.environ, NR_ENV):
                logger = NewRelicLogger()
        assert logger.enabled is False

    def test_record_content_default_true(self):
        logger = make_logger()
        assert logger.record_content is True

    def test_record_content_disabled_by_param(self):
        logger = make_logger(turn_off_message_logging=True)
        assert logger.record_content is False

    def test_record_content_disabled_by_env_var(self):
        with patch("newrelic.agent.register_application"):
            with patch.dict(
                os.environ,
                {**NR_ENV, "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": "false"},
            ):
                logger = NewRelicLogger()
                assert logger.record_content is False

    def test_record_content_requires_both_enabled(self):
        """param says record, but env var says no — result is False."""
        with patch("newrelic.agent.register_application"):
            with patch.dict(
                os.environ,
                {**NR_ENV, "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": "false"},
            ):
                logger = NewRelicLogger(turn_off_message_logging=False)
                assert logger.record_content is False

    def test_constructor_kwargs_take_priority_over_global_params(self):
        """Constructor turn_off_message_logging=True must not be overwritten by
        litellm.newrelic_params which defaults turn_off_message_logging to False."""
        from litellm.types.integrations.newrelic import NewRelicInitParams

        with patch("newrelic.agent.register_application"):
            with patch.dict(os.environ, NR_ENV):
                with patch(
                    "litellm.newrelic_params",
                    NewRelicInitParams(turn_off_message_logging=False),
                ):
                    logger = NewRelicLogger(turn_off_message_logging=True)
        assert logger.record_content is False

    def test_newrelic_params_plain_dict_branch(self):
        """litellm.newrelic_params can be a plain dict; it should be validated
        through NewRelicInitParams and its values applied to the logger."""
        with patch("newrelic.agent.register_application"):
            with patch.dict(os.environ, NR_ENV):
                with patch(
                    "litellm.newrelic_params",
                    {"turn_off_message_logging": True},
                ):
                    logger = NewRelicLogger()
        assert logger.turn_off_message_logging is True


# ---------------------------------------------------------------------------
# _parse_bool_env
# ---------------------------------------------------------------------------


class TestParseBoolEnv:
    def setup_method(self):
        self.logger = make_logger()

    @pytest.mark.parametrize("raw", ["true", "TRUE", "True", "1", "yes", "on", "ON"])
    def test_truthy_values(self, raw):
        with patch.dict(os.environ, {"MY_VAR": raw}):
            assert self.logger._parse_bool_env("MY_VAR") is True

    @pytest.mark.parametrize("raw", ["false", "FALSE", "0", "no", "off", "Off"])
    def test_falsy_values(self, raw):
        with patch.dict(os.environ, {"MY_VAR": raw}):
            assert self.logger._parse_bool_env("MY_VAR") is False

    @pytest.mark.parametrize("raw", [" true ", "  1\t", "\nyes"])
    def test_whitespace_tolerance_truthy(self, raw):
        with patch.dict(os.environ, {"MY_VAR": raw}):
            assert self.logger._parse_bool_env("MY_VAR") is True

    @pytest.mark.parametrize("raw", [" false ", "  0\t", "\nno"])
    def test_whitespace_tolerance_falsy(self, raw):
        with patch.dict(os.environ, {"MY_VAR": raw}):
            assert self.logger._parse_bool_env("MY_VAR") is False

    def test_missing_uses_default(self):
        with patch.dict(os.environ, {}, clear=True):
            assert self.logger._parse_bool_env("MY_VAR", default=True) is True
            assert self.logger._parse_bool_env("MY_VAR", default=False) is False

    def test_empty_string_uses_default(self):
        with patch.dict(os.environ, {"MY_VAR": ""}):
            assert self.logger._parse_bool_env("MY_VAR", default=True) is True
            assert self.logger._parse_bool_env("MY_VAR", default=False) is False

    @pytest.mark.parametrize("raw", ["maybe", "2", "enabled", "tru"])
    def test_unrecognised_value_falls_back_to_default_with_warning(self, raw):
        with (
            patch.dict(os.environ, {"MY_VAR": raw}),
            patch.object(nr_module.verbose_logger, "warning") as mock_warn,
        ):
            assert self.logger._parse_bool_env("MY_VAR", default=True) is True
            assert self.logger._parse_bool_env("MY_VAR", default=False) is False
            assert mock_warn.call_count == 2
            # Warning should mention the variable name and the raw value
            for call in mock_warn.call_args_list:
                rendered = _rendered_log_message(call)
                assert "MY_VAR" in rendered
                assert repr(raw) in rendered


# ---------------------------------------------------------------------------
# _get_trace_context
# ---------------------------------------------------------------------------


class TestGetTraceContext:
    def setup_method(self):
        self.logger = make_logger()

    def test_extracts_trace_id_from_traceparent(self):
        kwargs = make_kwargs(
            traceparent="00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
        )
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"

    def test_generates_uuid_when_no_headers(self):
        kwargs = make_kwargs()
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id is not None
        assert (
            len(trace_id) == 32
        )  # 32-char lowercase hex, matches W3C traceparent format

    def test_generates_uuid_when_traceparent_malformed(self):
        kwargs = make_kwargs(traceparent="not-valid")
        trace_id = self.logger._get_trace_context(kwargs)
        # Falls back to a 32-char lowercase hex, matching W3C traceparent format
        assert trace_id is not None
        assert len(trace_id) == 32

    def test_extracts_trace_id_from_mixed_case_traceparent_header(self):
        # Callers passing headers directly may not normalise case; per W3C spec
        # header names are case-insensitive, so "Traceparent" must work too.
        kwargs = make_kwargs()
        kwargs["litellm_params"]["metadata"]["headers"] = {
            "Traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
        }
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"

    def test_parse_failure_falls_through_to_synthetic_uuid(self):
        """When parsing upstream sources raises, emit a synthetic UUID rather
        than dropping the event. NR schema requires every AIM event carry a
        trace_id; this method's contract is to always return a valid string.
        """
        # Non-dict headers value forces .items() to raise inside the try
        kwargs = {"litellm_params": {"metadata": {"headers": "not-a-dict"}}}
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id is not None
        assert len(trace_id) == 32  # 32-char lowercase hex fallback


# ---------------------------------------------------------------------------
# _extract_message_content edge cases
# ---------------------------------------------------------------------------


class TestExtractMessageContent:
    def setup_method(self):
        self.logger = make_logger()

    def test_plain_text(self):
        assert self.logger._extract_message_content({"content": "hello"}) == "hello"

    def test_none_content_returns_empty_string(self):
        assert self.logger._extract_message_content({"content": None}) == ""

    def test_missing_content_returns_empty_string(self):
        assert self.logger._extract_message_content({}) == ""

    def test_tool_calls_serialized_as_json(self):
        msg = {
            "content": None,
            "tool_calls": [{"id": "call_1", "function": {"name": "get_weather"}}],
        }
        result = self.logger._extract_message_content(msg)
        assert "get_weather" in result
        assert "call_1" in result

    def test_multimodal_list_serialized_as_json(self):
        msg = {
            "content": [
                {"type": "text", "text": "describe this"},
                {"type": "image_url"},
            ]
        }
        result = self.logger._extract_message_content(msg)
        assert "describe this" in result
        assert "image_url" in result

    def test_non_string_content_coerced_to_str(self):
        """Numeric/bool content passes the None and list guards; final branch coerces to str."""
        assert self.logger._extract_message_content({"content": 123}) == "123"
        assert self.logger._extract_message_content({"content": True}) == "True"


# ---------------------------------------------------------------------------
# _extract_all_messages — record_content=False path
# ---------------------------------------------------------------------------


class TestExtractAllMessagesContentDisabled:
    def test_no_content_key_when_recording_disabled(self):
        logger = make_logger(turn_off_message_logging=True)
        kwargs = make_kwargs(messages=[{"role": "user", "content": "secret"}])
        response = make_response(content="also secret")

        messages = logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        for msg in messages:
            assert "content" not in msg


class TestExtractAllMessagesRespectsLitellmRedaction:
    """Regression tests for the async-streaming redaction bypass.

    NR-specific switches alone are insufficient: when
    ``litellm.turn_off_message_logging=True`` (or the per-request equivalents),
    async streaming callbacks receive an unredacted
    ``async_complete_streaming_response``. Without consulting LiteLLM's
    redaction decision the integration would still write generated content
    into NR events.
    """

    def _assert_no_content(self, logger, kwargs):
        response = make_response(content="streamed assistant text")

        messages = logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        # All extracted messages must carry no content payload
        for msg in messages:
            assert (
                "content" not in msg
            ), f"content leaked despite redaction signal: {msg}"
        # And there must actually be at least one user + one assistant entry,
        # otherwise the test would pass vacuously.
        assert any(not m.get("is_response") for m in messages)
        assert any(m.get("is_response") for m in messages)

    def test_global_turn_off_message_logging_blocks_content(self, monkeypatch):
        monkeypatch.setattr(litellm, "turn_off_message_logging", True)
        logger = make_logger()
        assert logger.record_content is True

        kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}])
        self._assert_no_content(logger, kwargs)

    def test_dynamic_param_turn_off_message_logging_blocks_content(self):
        logger = make_logger()
        assert logger.record_content is True

        kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}])
        kwargs["standard_callback_dynamic_params"] = {
            "turn_off_message_logging": True,
        }
        self._assert_no_content(logger, kwargs)

    def test_enable_redaction_header_blocks_content(self):
        logger = make_logger()
        assert logger.record_content is True

        kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}])
        kwargs["litellm_params"]["metadata"]["headers"] = {
            "x-litellm-enable-message-redaction": True,
        }
        self._assert_no_content(logger, kwargs)

    def test_dynamic_param_explicit_false_overrides_global_redaction(self, monkeypatch):
        """The dynamic param has higher priority than the global flag (see
        should_redact_message_logging). When a caller explicitly opts back into
        message logging per-request, NR must record content again."""
        monkeypatch.setattr(litellm, "turn_off_message_logging", True)
        logger = make_logger()

        kwargs = make_kwargs(messages=[{"role": "user", "content": "ok to log"}])
        kwargs["standard_callback_dynamic_params"] = {
            "turn_off_message_logging": False,
        }
        response = make_response(content="response text")

        messages = logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        request_msg = next(m for m in messages if not m.get("is_response"))
        response_msg = next(m for m in messages if m.get("is_response"))
        assert request_msg["content"] == "ok to log"
        assert response_msg["content"] == "response text"


class TestExtractAllMessagesTimestamps:
    def setup_method(self):
        self.logger = make_logger()

    def test_input_messages_get_start_time_timestamp(self):
        kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
        # make_kwargs sets start_time=1_000_000.0 and end_time=1_000_001.5
        response = make_response()

        messages = self.logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        input_msg = next(m for m in messages if not m.get("is_response"))
        assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0)

    def test_output_messages_get_end_time_timestamp(self):
        kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
        response = make_response()

        messages = self.logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        output_msg = next(m for m in messages if m.get("is_response"))
        assert output_msg["timestamp"] == int(1_000_001.5 * 1000.0)

    def test_timestamp_forwarded_to_event_data(self):
        logger = make_logger()
        mock_app = MagicMock()
        mock_app.enabled = True

        kwargs = make_kwargs(
            traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01",
            messages=[{"role": "user", "content": "Hi"}],
        )
        response = make_response()

        with patch("newrelic.agent.application", return_value=mock_app):
            logger._process_success(kwargs, response, start_time=1.0, end_time=2.5)

        calls = mock_app.record_custom_event.call_args_list
        message_events = [
            c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage"
        ]
        for event in message_events:
            assert "timestamp" in event


# ---------------------------------------------------------------------------
# Streaming response handling
# ---------------------------------------------------------------------------


def make_streaming_response(
    model="gpt-4",
    response_id="chatcmpl-stream123",
    content="Hello from streaming!",
    finish_reason="stop",
    prompt_tokens=8,
    completion_tokens=15,
):
    """Build a streaming-assembled response dict using 'delta' instead of 'message'."""
    return {
        "id": response_id,
        "model": model,
        "choices": [
            {
                "delta": {"role": "assistant", "content": content},
                "finish_reason": finish_reason,
            }
        ],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }


class TestStreamingResponse:
    """Verify graceful handling of streaming-assembled responses.

    When LiteLLM assembles a streaming response, some providers produce a
    final choice dict with a 'delta' key instead of 'message'. The integration
    must extract content from either key without raising.
    """

    def setup_method(self):
        self.logger = make_logger()

    def test_extracts_content_from_delta_key(self):
        kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
        response = make_streaming_response(content="Streamed reply")

        messages = self.logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        response_msgs = [m for m in messages if m.get("is_response")]
        assert len(response_msgs) == 1
        assert response_msgs[0]["content"] == "Streamed reply"
        assert response_msgs[0]["role"] == "assistant"

    def test_streaming_response_records_summary_and_message_events(self):
        mock_app = MagicMock()
        mock_app.enabled = True

        kwargs = make_kwargs(
            traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01",
            messages=[{"role": "user", "content": "Hi"}],
        )
        response = make_streaming_response(
            response_id="chatcmpl-stream123",
            content="Streamed reply",
            finish_reason="stop",
            prompt_tokens=8,
            completion_tokens=15,
        )

        with patch("newrelic.agent.application", return_value=mock_app):
            self.logger._process_success(kwargs, response, start_time=1.0, end_time=2.0)

        calls = mock_app.record_custom_event.call_args_list
        event_types = [c[0][0] for c in calls]
        assert "LlmChatCompletionSummary" in event_types
        assert "LlmChatCompletionMessage" in event_types

        message_events = [
            c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage"
        ]
        response_msg = next((e for e in message_events if e.get("is_response")), None)
        assert response_msg is not None
        assert response_msg["content"] == "Streamed reply"

    @pytest.mark.asyncio
    async def test_async_log_success_event_streaming(self):
        """async_log_success_event is the primary entry point for streaming calls."""
        mock_app = MagicMock()
        mock_app.enabled = True

        kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
        response = make_streaming_response()

        with patch("newrelic.agent.application", return_value=mock_app):
            await self.logger.async_log_success_event(
                kwargs, response, start_time=1.0, end_time=2.0
            )

        calls = mock_app.record_custom_event.call_args_list
        event_types = [c[0][0] for c in calls]
        assert "LlmChatCompletionSummary" in event_types
        assert "LlmChatCompletionMessage" in event_types

    def test_no_content_when_recording_disabled_streaming(self):
        logger = make_logger(turn_off_message_logging=True)
        kwargs = make_kwargs(messages=[{"role": "user", "content": "secret"}])
        response = make_streaming_response(content="also secret")

        messages = logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )

        for msg in messages:
            assert "content" not in msg


# ---------------------------------------------------------------------------
# Explicit-None defensive tests
# ---------------------------------------------------------------------------


class TestExplicitNoneValues:
    """Verify that explicitly None values in kwargs/response don't raise or silently drop events."""

    def setup_method(self):
        self.logger = make_logger()

    # _get_trace_context — chained dict lookups
    def test_trace_context_litellm_params_none(self):
        kwargs = make_kwargs()
        kwargs["litellm_params"] = None
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id is not None  # falls back to UUID

    def test_trace_context_metadata_none(self):
        kwargs = make_kwargs()
        kwargs["litellm_params"] = {"metadata": None}
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id is not None

    def test_trace_context_headers_none(self):
        kwargs = make_kwargs()
        kwargs["litellm_params"] = {"metadata": {"headers": None}}
        trace_id = self.logger._get_trace_context(kwargs)
        assert trace_id is not None

    # _get_request_params
    def test_request_params_optional_params_none(self):
        assert self.logger._get_request_params({"optional_params": None}) == {}

    # _get_model_names
    def test_model_names_model_none_in_kwargs(self):
        request_model, _ = self.logger._get_model_names(
            {"model": None}, make_response()
        )
        assert request_model == "unknown"

    def test_model_names_model_none_in_response(self):
        response = make_response()
        response["model"] = None
        _, response_model = self.logger._get_model_names(make_kwargs(), response)
        assert response_model == "gpt-4"  # falls back to request_model from kwargs

    # _extract_all_messages
    def test_extract_messages_messages_none(self):
        kwargs = make_kwargs()
        kwargs["messages"] = None
        response = make_response()
        messages = self.logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )
        # No request messages, but response message should still be extracted
        assert any(m.get("is_response") for m in messages)

    def test_extract_messages_choices_none(self):
        kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}])
        response = make_response()
        response["choices"] = None
        messages = self.logger._extract_all_messages(
            kwargs, response, response_model="gpt-4", vendor="openai"
        )
        # No response messages, but request message should still be extracted
        assert any(not m.get("is_response") for m in messages)


# ---------------------------------------------------------------------------
# Helper edge cases
# ---------------------------------------------------------------------------


class TestExtractUsage:
    def setup_method(self):
        self.logger = make_logger()

    def test_missing_usage_returns_zeros(self):
        response = {"id": "r1", "model": "gpt-4", "choices": []}
        usage = self.logger._extract_usage(response)
        assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    def test_explicit_none_token_fields_return_zeros(self):
        response = {
            "usage": {
                "prompt_tokens": None,
                "completion_tokens": None,
                "total_tokens": None,
            }
        }
        usage = self.logger._extract_usage(response)
        assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}


class TestGetFinishReason:
    def setup_method(self):
        self.logger = make_logger()

    def test_returns_unknown_when_no_choices(self):
        response = {"choices": []}
        assert self.logger._get_finish_reason(response) == "unknown"

    def test_returns_unknown_when_choices_missing(self):
        assert self.logger._get_finish_reason({}) == "unknown"

    def test_returns_unknown_when_finish_reason_explicitly_none(self):
        response = {"choices": [{"finish_reason": None}]}
        assert self.logger._get_finish_reason(response) == "unknown"


class TestToEpochMs:
    def setup_method(self):
        self.logger = make_logger()

    def test_float_passthrough(self):
        assert self.logger._to_epoch_ms(1.0) == pytest.approx(1000.0)

    def test_datetime_converted(self):
        dt = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
        assert self.logger._to_epoch_ms(dt) == pytest.approx(dt.timestamp() * 1000.0)


class TestGetDuration:
    def setup_method(self):
        self.logger = make_logger()

    def test_uses_kwargs_value_when_present(self):
        kwargs = {"llm_api_duration_ms": 750.0}
        assert self.logger._get_duration(kwargs, 0.0, 1.0) == 750.0

    def test_calculates_from_float_timestamps(self):
        kwargs = {}
        result = self.logger._get_duration(kwargs, 1.0, 2.5)
        assert result == pytest.approx(1500.0)

    def test_calculates_from_datetime_timestamps(self):
        kwargs = {}
        start = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
        end = datetime(2024, 1, 1, 0, 0, 1, 500000, tzinfo=timezone.utc)  # +1.5s
        result = self.logger._get_duration(kwargs, start, end)
        assert result == pytest.approx(1500.0)

    def test_returns_none_when_nothing_available(self):
        assert self.logger._get_duration({}, None, None) is None


class TestGetRequestParams:
    def setup_method(self):
        self.logger = make_logger()

    def test_includes_only_present_params(self):
        kwargs = {"optional_params": {"temperature": 0.7}}
        params = self.logger._get_request_params(kwargs)
        assert params == {"temperature": 0.7}
        assert "max_tokens" not in params

    def test_empty_when_no_optional_params(self):
        assert self.logger._get_request_params({}) == {}


# ---------------------------------------------------------------------------
# _process_success — comprehensive happy-path
# ---------------------------------------------------------------------------


class TestProcessSuccess:
    def test_records_summary_and_message_events(self):
        logger = make_logger()
        mock_app = MagicMock()
        mock_app.enabled = True

        kwargs = make_kwargs(
            traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01",
            messages=[{"role": "user", "content": "Hello"}],
            optional_params={"temperature": 0.5, "max_tokens": 100},
        )
        response = make_response(
            response_id="chatcmpl-xyz",
            content="Hi there!",
            finish_reason="stop",
            prompt_tokens=5,
            completion_tokens=10,
        )

        with patch("newrelic.agent.application", return_value=mock_app):
            logger._process_success(kwargs, response, start_time=1.0, end_time=2.5)

        calls = mock_app.record_custom_event.call_args_list
        event_types = [c[0][0] for c in calls]

        assert "LlmChatCompletionSummary" in event_types
        assert "LlmChatCompletionMessage" in event_types

        # Verify summary event fields
        summary_data = next(
            c[0][1] for c in calls if c[0][0] == "LlmChatCompletionSummary"
        )
        assert summary_data["vendor"] == "openai"
        assert summary_data["request.model"] == "gpt-4"
        assert summary_data["response.model"] == "gpt-4"
        assert summary_data["response.choices.finish_reason"] == "stop"
        assert summary_data["response.usage.prompt_tokens"] == 5
        assert summary_data["response.usage.completion_tokens"] == 10
        assert summary_data["response.usage.total_tokens"] == 15
        assert summary_data["request.temperature"] == 0.5
        assert summary_data["request.max_tokens"] == 100
        assert summary_data["ingest_source"] == "litellm"
        assert summary_data["trace_id"] == "aabbccddeeff00112233445566778899"

        # Verify message event id format: "{llm_response_id}-{sequence}"
        message_events = [
            c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage"
        ]
        assert any(e["id"].startswith("chatcmpl-xyz-") for e in message_events)
        response_msg = next(e for e in message_events if e.get("is_response"))
        assert response_msg["content"] == "Hi there!"
        assert response_msg["role"] == "assistant"

    def test_skips_when_disabled(self):
        logger = make_logger()
        logger.enabled = False

        with patch("newrelic.agent.application") as mock_app:
            logger._process_success(make_kwargs(), make_response())

        mock_app.assert_not_called()


# ---------------------------------------------------------------------------
# _record_error_metric
# ---------------------------------------------------------------------------


class TestRecordErrorMetric:
    def setup_method(self):
        self.logger = make_logger()

    def test_calls_record_custom_metric(self):
        mock_app = MagicMock()
        mock_app.enabled = True

        with patch.object(self.logger, "_check_and_emit_periodic_metric"):
            with patch("newrelic.agent.application", return_value=mock_app):
                self.logger._record_error_metric()

        mock_app.record_custom_metric.assert_called_once_with("LLM/LiteLLM/Error", 1)

    def test_skips_when_app_disabled(self):
        mock_app = MagicMock()
        mock_app.enabled = False

        with patch.object(self.logger, "_check_and_emit_periodic_metric"):
            with patch("newrelic.agent.application", return_value=mock_app):
                self.logger._record_error_metric()

        mock_app.record_custom_metric.assert_not_called()

    def test_calls_check_and_emit_periodic_metric(self):
        with patch.object(
            self.logger, "_check_and_emit_periodic_metric"
        ) as mock_periodic:
            with patch("newrelic.agent.application", return_value=MagicMock()):
                self.logger._record_error_metric()

        mock_periodic.assert_called_once()

    def test_skips_when_logger_disabled(self):
        self.logger.enabled = False
        with patch("newrelic.agent.application") as mock_app:
            self.logger._record_error_metric()
        mock_app.assert_not_called()

    def test_handles_exception(self):
        with patch(
            "newrelic.agent.application", side_effect=RuntimeError("agent down")
        ):
            self.logger._record_error_metric()  # must not raise


# ---------------------------------------------------------------------------
# _emit_supportability_metric
# ---------------------------------------------------------------------------


class TestEmitSupportabilityMetric:
    def setup_method(self):
        self.logger = make_logger()
        NewRelicLogger._last_metric_emission_time = 0.0

    def test_records_metric_with_correct_name_and_value(self):
        mock_app = MagicMock()
        mock_app.enabled = True
        with patch("newrelic.agent.application", return_value=mock_app):
            with patch.object(
                self.logger, "_get_litellm_version", return_value="1.80.0"
            ):
                self.logger._emit_supportability_metric()
        mock_app.record_custom_metric.assert_called_once_with(
            "Supportability/Python/ML/LiteLLM/1.80.0", 1
        )

    def test_updates_last_emission_time(self):
        mock_app = MagicMock()
        mock_app.enabled = True
        fake_now = 9_999_999.0
        with patch("newrelic.agent.application", return_value=mock_app):
            with patch(
                "litellm.integrations.newrelic.newrelic.time.time",
                return_value=fake_now,
            ):
                self.logger._emit_supportability_metric()
        assert NewRelicLogger._last_metric_emission_time == fake_now

    def test_skips_when_app_disabled(self):
        mock_app = MagicMock()
        mock_app.enabled = False
        with patch("newrelic.agent.application", return_value=mock_app):
            self.logger._emit_supportability_metric()
        mock_app.record_custom_metric.assert_not_called()
        # Timestamp is still updated to back off lock contention during registration.
        assert NewRelicLogger._last_metric_emission_time != 0.0

    def test_skips_when_no_app(self):
        with patch("newrelic.agent.application", return_value=None):
            self.logger._emit_supportability_metric()
        # Timestamp is updated even when app is None to back off lock contention
        # if the agent never starts or is slow to initialise.
        assert NewRelicLogger._last_metric_emission_time != 0.0

    def test_handles_exception(self):
        with patch(
            "newrelic.agent.application", side_effect=RuntimeError("agent down")
        ):
            self.logger._emit_supportability_metric()  # must not raise


# ---------------------------------------------------------------------------
# _check_and_emit_periodic_metric
# ---------------------------------------------------------------------------


class TestCheckAndEmitPeriodicMetric:
    def setup_method(self):
        self.logger = make_logger()
        NewRelicLogger._last_metric_emission_time = 0.0

    def test_emits_on_first_call(self):
        """_last_metric_emission_time starts at 0.0; any real time satisfies 27-hour window."""
        with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
            with patch(
                "litellm.integrations.newrelic.newrelic.time.time",
                return_value=100_000.0,
            ):
                self.logger._check_and_emit_periodic_metric()
        mock_emit.assert_called_once()

    def test_does_not_re_emit_within_27_hours(self):
        recent = 1_000_000.0
        NewRelicLogger._last_metric_emission_time = recent
        with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
            with patch(
                "litellm.integrations.newrelic.newrelic.time.time",
                return_value=recent + 3600,  # 1 hour later
            ):
                self.logger._check_and_emit_periodic_metric()
        mock_emit.assert_not_called()

    def test_re_emits_after_27_hours(self):
        old = 1_000_000.0
        NewRelicLogger._last_metric_emission_time = old
        with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
            with patch(
                "litellm.integrations.newrelic.newrelic.time.time",
                return_value=old + 97201,  # 27 hours + 1 second
            ):
                self.logger._check_and_emit_periodic_metric()
        mock_emit.assert_called_once()

    def test_boundary_exactly_27_hours_triggers_emission(self):
        old = 1_000_000.0
        NewRelicLogger._last_metric_emission_time = old
        with patch.object(self.logger, "_emit_supportability_metric") as mock_emit:
            with patch(
                "litellm.integrations.newrelic.newrelic.time.time",
                return_value=old + 97200,
            ):
                self.logger._check_and_emit_periodic_metric()
        mock_emit.assert_called_once()


# ---------------------------------------------------------------------------
# _get_litellm_version
# ---------------------------------------------------------------------------


class TestGetLitellmVersion:
    def setup_method(self):
        self.logger = make_logger()

    def test_returns_unknown_on_exception(self):
        with patch("importlib.metadata.version", side_effect=Exception("no package")):
            result = self.logger._get_litellm_version()
        assert result == "unknown"


# ---------------------------------------------------------------------------
# _record_summary_event — disabled-app and exception paths
# ---------------------------------------------------------------------------

_USAGE = {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}


class TestRecordSummaryEvent:
    def setup_method(self):
        self.logger = make_logger()

    def _call(self, **kwargs):
        self.logger._record_summary_event(
            request_id="req-1",
            trace_id="trace-abc",
            request_model="gpt-4",
            response_model="gpt-4",
            vendor="openai",
            finish_reason="stop",
            num_messages=2,
            usage=_USAGE,
            **kwargs,
        )

    def test_skips_when_app_disabled(self):
        mock_app = MagicMock()
        mock_app.enabled = False
        with patch("newrelic.agent.application", return_value=mock_app):
            self._call()
        mock_app.record_custom_event.assert_not_called()

    def test_handles_exception(self):
        with patch(
            "newrelic.agent.application", side_effect=RuntimeError("agent down")
        ):
            self._call()  # must not raise


# ---------------------------------------------------------------------------
# _record_message_events — disabled-app and exception paths
# ---------------------------------------------------------------------------

_MESSAGES = [
    {"role": "user", "sequence": 0, "response.model": "gpt-4", "vendor": "openai"}
]


class TestRecordMessageEvents:
    def setup_method(self):
        self.logger = make_logger()

    def _call(self):
        self.logger._record_message_events(
            request_id="req-1",
            llm_response_id="resp-1",
            trace_id="trace-abc",
            messages=_MESSAGES,
        )

    def test_skips_when_app_disabled(self):
        mock_app = MagicMock()
        mock_app.enabled = False
        with patch("newrelic.agent.application", return_value=mock_app):
            self._call()
        mock_app.record_custom_event.assert_not_called()

    def test_handles_exception(self):
        with patch(
            "newrelic.agent.application", side_effect=RuntimeError("agent down")
        ):
            self._call()  # must not raise


# ---------------------------------------------------------------------------
# CustomLogger interface entry points
# ---------------------------------------------------------------------------


class TestLogSuccessEvent:
    def test_delegates_to_process_success(self):
        logger = make_logger()
        with patch.object(logger, "_process_success") as mock_process:
            logger.log_success_event(make_kwargs(), make_response(), 1.0, 2.0)
        mock_process.assert_called_once()

    def test_exception_is_handled(self):
        logger = make_logger()
        with patch.object(logger, "_process_success", side_effect=RuntimeError("boom")):
            logger.log_success_event(make_kwargs(), make_response(), 1.0, 2.0)

    @pytest.mark.asyncio
    async def test_async_delegates_to_process_success(self):
        logger = make_logger()
        with patch.object(logger, "_process_success") as mock_process:
            await logger.async_log_success_event(
                make_kwargs(), make_response(), 1.0, 2.0
            )
        mock_process.assert_called_once()

    @pytest.mark.asyncio
    async def test_async_exception_is_handled(self):
        logger = make_logger()
        with patch.object(logger, "_process_success", side_effect=RuntimeError("boom")):
            await logger.async_log_success_event(
                make_kwargs(), make_response(), 1.0, 2.0
            )


class TestLogFailureEvent:
    def test_sync_records_error_metric(self):
        logger = make_logger()
        with patch.object(logger, "_record_error_metric") as mock_metric:
            logger.log_failure_event(make_kwargs(), None, 1.0, 2.0)
        mock_metric.assert_called_once()

    def test_sync_exception_is_handled(self):
        logger = make_logger()
        with patch.object(
            logger, "_record_error_metric", side_effect=RuntimeError("boom")
        ):
            logger.log_failure_event(make_kwargs(), None, 1.0, 2.0)

    @pytest.mark.asyncio
    async def test_async_records_error_metric(self):
        logger = make_logger()
        with patch.object(logger, "_record_error_metric") as mock_metric:
            await logger.async_log_failure_event(make_kwargs(), None, 1.0, 2.0)
        mock_metric.assert_called_once()

    @pytest.mark.asyncio
    async def test_async_exception_is_handled(self):
        logger = make_logger()
        with patch.object(
            logger, "_record_error_metric", side_effect=RuntimeError("boom")
        ):
            await logger.async_log_failure_event(make_kwargs(), None, 1.0, 2.0)


# ---------------------------------------------------------------------------
# async_health_check
# ---------------------------------------------------------------------------


class TestAsyncHealthCheck:
    @pytest.mark.asyncio
    async def test_unhealthy_when_disabled(self):
        logger = make_logger()
        logger.enabled = False
        result = await logger.async_health_check()
        assert result["status"] == "unhealthy"
        assert result["error_message"] is not None

    @pytest.mark.asyncio
    async def test_healthy_when_app_enabled_records_test_event(self):
        logger = make_logger()
        mock_app = MagicMock()
        mock_app.enabled = True
        with patch("newrelic.agent.application", return_value=mock_app):
            result = await logger.async_health_check()
        assert result["status"] == "healthy"
        assert result["error_message"] is None

        mock_app.record_custom_event.assert_called_once()
        event_type, event_data = mock_app.record_custom_event.call_args[0]
        assert event_type == "LiteLLMConnectionTest"
        assert event_data["is_test_event"] is True
        assert event_data["app_name"] == logger.app_name
        assert event_data["source"] == "litellm-proxy"
        assert isinstance(event_data["timestamp"], float)

    @pytest.mark.asyncio
    async def test_unhealthy_when_app_disabled(self):
        logger = make_logger()
        mock_app = MagicMock()
        mock_app.enabled = False
        with patch("newrelic.agent.application", return_value=mock_app):
            result = await logger.async_health_check()
        assert result["status"] == "unhealthy"
        assert result["error_message"] is not None
        mock_app.record_custom_event.assert_not_called()

    @pytest.mark.asyncio
    async def test_exception_returns_unhealthy(self):
        logger = make_logger()
        with patch(
            "newrelic.agent.application", side_effect=RuntimeError("agent down")
        ):
            result = await logger.async_health_check()
        assert result["status"] == "unhealthy"
        assert "agent down" in result["error_message"]

    @pytest.mark.asyncio
    async def test_record_custom_event_failure_returns_unhealthy(self):
        logger = make_logger()
        mock_app = MagicMock()
        mock_app.enabled = True
        mock_app.record_custom_event.side_effect = RuntimeError("intake unreachable")
        with patch("newrelic.agent.application", return_value=mock_app):
            result = await logger.async_health_check()
        assert result["status"] == "unhealthy"
        assert "intake unreachable" in result["error_message"]


# ---------------------------------------------------------------------------
# _extract_completion_id fallback chain
# ---------------------------------------------------------------------------


class TestExtractCompletionId:
    def setup_method(self):
        self.logger = make_logger()

    def test_uses_litellm_call_id_when_response_has_no_id(self):
        result = self.logger._extract_completion_id(
            kwargs={"litellm_call_id": "call-abc-123"},
            response_obj={},
        )
        assert result == "call-abc-123"

    def test_generates_uuid_when_neither_id_present(self):
        result = self.logger._extract_completion_id(kwargs={}, response_obj={})
        # UUID4 hex-with-dashes is 36 chars; just confirm shape and uniqueness
        assert isinstance(result, str)
        assert len(result) == 36
        second = self.logger._extract_completion_id(kwargs={}, response_obj={})
        assert result != second


# ---------------------------------------------------------------------------
# StandardLoggingPayload preference across extractors
# ---------------------------------------------------------------------------


class TestStandardLoggingPayloadPreference:
    """Each extractor that accepts a StandardLoggingPayload must prefer its
    values over the raw kwargs/response fallbacks."""

    def setup_method(self):
        self.logger = make_logger()

    def test_trace_context_uses_slo_trace_id_when_no_traceparent(self):
        kwargs = {"litellm_params": {"metadata": {"headers": {}}}}
        trace_id = self.logger._get_trace_context(
            kwargs, standard_logging_object=make_slo()
        )
        assert trace_id == "slo-trace-abc"

    def test_vendor_from_slo(self):
        # kwargs carries a different provider; SLO must win.
        kwargs = {"litellm_params": {"custom_llm_provider": "kwargs-provider"}}
        assert (
            self.logger._get_vendor(kwargs, standard_logging_object=make_slo())
            == "slo-provider"
        )

    def test_model_names_uses_slo_model(self):
        request_model, _ = self.logger._get_model_names(
            {"model": "kwargs-model"},
            make_response(model="response-model"),
            standard_logging_object=make_slo(),
        )
        assert request_model == "slo-model"

    def test_usage_from_slo_when_any_token_field_present(self):
        # make_response defaults to 10/20/30 tokens; SLO sentinels are 100/200/300.
        usage = self.logger._extract_usage(
            make_response(), standard_logging_object=make_slo()
        )
        assert usage == {
            "prompt_tokens": 100,
            "completion_tokens": 200,
            "total_tokens": 300,
        }

    def test_duration_from_slo_response_time_converted_to_ms(self):
        # SLO response_time is 1.5 seconds; expected 1500.0 ms.
        # Pass start/end that would compute a different value to prove SLO won.
        duration = self.logger._get_duration(
            kwargs={"llm_api_duration_ms": 9999.0},
            start_time=1.0,
            end_time=2.0,
            standard_logging_object=make_slo(),
        )
        assert duration == 1500.0

    def test_request_params_from_slo_model_parameters(self):
        params = self.logger._get_request_params(
            {"optional_params": {"temperature": 0.1}},
            standard_logging_object=make_slo(),
        )
        assert params == {"temperature": 0.7, "max_tokens": 500}

    def test_extract_all_messages_sources_timestamps_and_messages_from_slo(self):
        """Covers three SLO branches at once: startTime, endTime, and messages list."""
        kwargs = make_kwargs(messages=[{"role": "user", "content": "from-kwargs"}])
        messages = self.logger._extract_all_messages(
            kwargs,
            make_response(),
            response_model="gpt-4",
            vendor="openai",
            standard_logging_object=make_slo(),
        )

        request = next(m for m in messages if not m.get("is_response"))
        assert request["content"] == "from-slo"  # SLO messages list wins
        assert request["timestamp"] == int(2_000_000.0 * 1000.0)  # SLO startTime

        response = next(m for m in messages if m.get("is_response"))
        assert response["timestamp"] == int(2_000_001.5 * 1000.0)  # SLO endTime
