"""
Tests for max_tokens vs thinking.budget_tokens constraint validation
in the websearch interception agentic loop.

Covers:
  - M1-I1: max_tokens auto-adjustment when <= thinking.budget_tokens
  - M1-I3: Unit tests for thinking parameter validation
  - M2-I5/I8: litellm_logging_obj excluded from follow-up kwargs to prevent SpendLog dedup
  - M3-I12: Regression tests for error scenarios
"""

from typing import Any, Dict, List
from unittest.mock import MagicMock, patch

import pytest

from litellm.integrations.websearch_interception.handler import (
    WebSearchInterceptionLogger,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_tool_calls() -> List[Dict]:
    return [
        {
            "id": "toolu_01",
            "type": "tool_use",
            "name": "web_search",
            "input": {"query": "litellm spend tracking"},
        }
    ]


def _make_logging_obj(
    model: str = "bedrock/us.anthropic.claude-opus-4-6-v1",
) -> MagicMock:
    obj = MagicMock()
    obj.model_call_details = {
        "agentic_loop_params": {"model": model, "custom_llm_provider": "bedrock"},
    }
    return obj


# ---------------------------------------------------------------------------
# M1-I1 / M1-I3: max_tokens validation against thinking.budget_tokens
# ---------------------------------------------------------------------------


class TestThinkingBudgetTokensConstraint:
    """Validate that _execute_agentic_loop adjusts max_tokens when <= thinking.budget_tokens."""

    @pytest.mark.asyncio
    async def test_max_tokens_adjusted_when_less_than_budget(self):
        """max_tokens < thinking.budget_tokens → auto-adjusted to budget_tokens + 1024."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()  # dummy response

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={
                    "max_tokens": 1024,
                    "thinking": {"type": "enabled", "budget_tokens": 5000},
                },
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={},
            )

        assert captured_kwargs["max_tokens"] == 5000 + 1024

    @pytest.mark.asyncio
    async def test_max_tokens_adjusted_when_equal_to_budget(self):
        """max_tokens == thinking.budget_tokens → still adjusted (must be strictly greater)."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={
                    "max_tokens": 5000,
                    "thinking": {"type": "enabled", "budget_tokens": 5000},
                },
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={},
            )

        assert captured_kwargs["max_tokens"] == 5000 + 1024

    @pytest.mark.asyncio
    async def test_max_tokens_unchanged_when_greater_than_budget(self):
        """max_tokens > thinking.budget_tokens → no adjustment needed."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={
                    "max_tokens": 10000,
                    "thinking": {"type": "enabled", "budget_tokens": 5000},
                },
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={},
            )

        assert captured_kwargs["max_tokens"] == 10000

    @pytest.mark.asyncio
    async def test_no_thinking_param_no_adjustment(self):
        """No thinking parameter → max_tokens used as-is (default 1024)."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={},
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={},
            )

        assert captured_kwargs["max_tokens"] == 1024

    @pytest.mark.asyncio
    async def test_thinking_without_budget_tokens_no_adjustment(self):
        """thinking param exists but has no budget_tokens → max_tokens used as-is."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={
                    "max_tokens": 2048,
                    "thinking": {"type": "enabled"},
                },
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={},
            )

        assert captured_kwargs["max_tokens"] == 2048


class TestResolveMaxTokensEdgeCases:
    """Edge cases for _resolve_max_tokens: infinity, negative, extreme values."""

    def test_infinity_budget_tokens_no_crash(self):
        """float('inf') budget_tokens must not crash with OverflowError."""
        result = WebSearchInterceptionLogger._resolve_max_tokens(
            {"max_tokens": 1024, "thinking": {"budget_tokens": float("inf")}}, {}
        )
        assert result == 1024  # no adjustment for non-finite values

    def test_negative_infinity_no_crash(self):
        result = WebSearchInterceptionLogger._resolve_max_tokens(
            {"max_tokens": 1024, "thinking": {"budget_tokens": float("-inf")}}, {}
        )
        assert result == 1024

    def test_nan_budget_tokens_no_crash(self):
        result = WebSearchInterceptionLogger._resolve_max_tokens(
            {"max_tokens": 1024, "thinking": {"budget_tokens": float("nan")}}, {}
        )
        assert result == 1024

    def test_negative_budget_tokens_no_adjustment(self):
        result = WebSearchInterceptionLogger._resolve_max_tokens(
            {"max_tokens": 1024, "thinking": {"budget_tokens": -100}}, {}
        )
        assert result == 1024

    def test_zero_budget_tokens_no_adjustment(self):
        result = WebSearchInterceptionLogger._resolve_max_tokens(
            {"max_tokens": 1024, "thinking": {"budget_tokens": 0}}, {}
        )
        assert result == 1024


# ---------------------------------------------------------------------------
# M2-I5 / M2-I8: litellm_logging_obj excluded from follow-up kwargs
# ---------------------------------------------------------------------------


class TestLoggingObjExcludedFromFollowUp:
    """Verify litellm_logging_obj is NOT forwarded to the follow-up acreate() call.

    Passing the same logging object to both initial and follow-up calls causes
    the has_logged_async_success dedup flag to fire, silently preventing the
    initial call's spend from being recorded in SpendLogs.
    """

    @pytest.mark.asyncio
    async def test_litellm_logging_obj_excluded_from_anthropic_followup(self):
        """The Anthropic messages follow-up must NOT receive litellm_logging_obj."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        fake_logging_obj = _make_logging_obj()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={"max_tokens": 4096},
                logging_obj=fake_logging_obj,
                stream=False,
                kwargs={
                    "litellm_logging_obj": fake_logging_obj,
                    "metadata": {"user_api_key": "test-key-hash"},
                    "temperature": 0.5,
                },
            )

        # litellm_logging_obj must be absent from the follow-up call
        assert "litellm_logging_obj" not in captured_kwargs
        # But other kwargs (metadata, temperature) must be preserved
        assert captured_kwargs.get("metadata") == {"user_api_key": "test-key-hash"}
        assert captured_kwargs.get("temperature") == 0.5

    @pytest.mark.asyncio
    async def test_websearch_flags_also_excluded(self):
        """Both _websearch_interception flags and litellm_logging_obj must be excluded."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={"max_tokens": 4096},
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={
                    "litellm_logging_obj": MagicMock(),
                    "_websearch_interception_converted_stream": True,
                    "_websearch_interception_other": "x",
                    "api_key": "fake",
                },
            )

        assert "litellm_logging_obj" not in captured_kwargs
        assert "_websearch_interception_converted_stream" not in captured_kwargs
        assert "_websearch_interception_other" not in captured_kwargs
        assert captured_kwargs.get("api_key") == "fake"


# ---------------------------------------------------------------------------
# M3-I12: Regression tests for error scenarios
# ---------------------------------------------------------------------------


class TestFollowUpErrorScenarios:
    """Regression tests: the agentic loop must surface errors properly and
    not silently swallow them (except at the _call_agentic_completion_hooks
    level which intentionally catches to return the initial response)."""

    @pytest.mark.asyncio
    async def test_followup_400_raises(self):
        """A 400 error from the follow-up call must propagate out of _execute_agentic_loop."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])

        async def _fail_acreate(**kw):
            raise Exception("max_tokens must be greater than thinking.budget_tokens")

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fail_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            with pytest.raises(Exception, match="max_tokens must be greater"):
                await logger._execute_agentic_loop(
                    model="us.anthropic.claude-opus-4-6-v1",
                    messages=[{"role": "user", "content": "hi"}],
                    tool_calls=_make_tool_calls(),
                    thinking_blocks=[],
                    anthropic_messages_optional_request_params={"max_tokens": 4096},
                    logging_obj=_make_logging_obj(),
                    stream=False,
                    kwargs={},
                )

    @pytest.mark.asyncio
    async def test_search_failure_does_not_crash_loop(self):
        """If a search fails, the loop should still attempt the follow-up with error text."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", side_effect=Exception("search API down")
            ),
        ):

            result = await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={"max_tokens": 4096},
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={},
            )

        # The follow-up call should have been made (with error text in search results)
        assert result is not None
        # Messages should contain the error text
        follow_up_messages = captured_kwargs.get("messages", [])
        assert len(follow_up_messages) > 1  # original + assistant + tool_result

    @pytest.mark.asyncio
    async def test_metadata_preserved_after_logging_obj_exclusion(self):
        """Proxy metadata (user_api_key, team_id, etc.) must survive in follow-up kwargs
        even after litellm_logging_obj is excluded — so the new logging_obj from
        function_setup has access to proxy tracking metadata."""
        logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
        captured_kwargs: Dict[str, Any] = {}

        async def _fake_acreate(**kw):
            captured_kwargs.update(kw)
            return MagicMock()

        proxy_metadata = {
            "user_api_key": "test-proxy-key-hash",
            "user_api_key_user_id": "user-123",
            "user_api_key_team_id": "team-456",
            "user_api_key_org_id": "org-789",
            "user_api_key_end_user_id": "end-user-001",
        }

        with (
            patch(
                "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
                side_effect=_fake_acreate,
            ),
            patch.object(
                logger, "_execute_search", return_value=("search result", None)
            ),
        ):

            await logger._execute_agentic_loop(
                model="us.anthropic.claude-opus-4-6-v1",
                messages=[{"role": "user", "content": "hi"}],
                tool_calls=_make_tool_calls(),
                thinking_blocks=[],
                anthropic_messages_optional_request_params={"max_tokens": 4096},
                logging_obj=_make_logging_obj(),
                stream=False,
                kwargs={
                    "litellm_logging_obj": MagicMock(),
                    "metadata": proxy_metadata,
                    "litellm_call_id": "call-abc-123",
                },
            )

        # litellm_logging_obj excluded
        assert "litellm_logging_obj" not in captured_kwargs
        # But ALL proxy metadata must be preserved
        assert captured_kwargs.get("metadata") == proxy_metadata
        assert captured_kwargs.get("litellm_call_id") == "call-abc-123"
