"""
Integration tests for Vertex AI rerank functionality.
These tests demonstrate end-to-end usage of the Vertex AI rerank feature.
"""

import importlib
from unittest.mock import MagicMock

import httpx


class TestVertexAIRerankIntegration:
    def setup_method(self):
        # Reload modules to ensure fresh references after conftest reloads litellm.
        # This ensures the class being patched is the same one used by the tests.
        import litellm.llms.vertex_ai.rerank.transformation as rerank_transformation_module

        importlib.reload(rerank_transformation_module)

        # Re-import after reload to get the fresh class
        from litellm.llms.vertex_ai.rerank.transformation import (
            VertexAIRerankConfig as FreshConfig,
        )

        self.config = FreshConfig()
        self.model = "semantic-ranker-default@latest"

    def test_end_to_end_rerank_flow(self):
        """
        Test complete rerank flow from request to response.

        Uses instance-level mocking to avoid class-reference issues caused by
        importlib.reload(litellm) in conftest.py.
        """
        # Mock authentication at instance level
        mock_ensure_access_token = MagicMock(
            return_value=("test-access-token", "test-project-123")
        )
        self.config._ensure_access_token = mock_ensure_access_token

        # Test documents
        documents = [
            "Gemini is a cutting edge large language model created by Google.",
            "The Gemini zodiac symbol often depicts two figures standing side-by-side.",
            "Gemini is a constellation that can be seen in the night sky.",
            "Google's Gemini AI model represents a significant advancement in artificial intelligence technology.",
        ]
        query = "What is Google Gemini?"

        # Step 1: Test request transformation
        # Validate environment
        headers = self.config.validate_environment(
            headers={}, model=self.model, api_key=None
        )

        # Transform request
        request_data = self.config.transform_rerank_request(
            model=self.model,
            optional_rerank_params={
                "query": query,
                "documents": documents,
                "top_n": 2,
                "return_documents": True,
            },
            headers=headers,
        )

        # Verify request structure
        assert request_data["model"] == self.model
        assert request_data["query"] == query
        assert request_data["topN"] == 2
        assert request_data["ignoreRecordDetailsInResponse"] == False
        assert len(request_data["records"]) == 4

        # Verify record structure
        for i, record in enumerate(request_data["records"]):
            assert record["id"] == str(i)  # 0-based indexing
            assert "title" in record
            assert "content" in record
            assert record["content"] == documents[i]

        # Step 2: Test response transformation
        # Mock Vertex AI Discovery Engine response
        mock_response_data = {
            "records": [
                {
                    "id": "3",
                    "score": 0.95,
                    "title": "Google's Gemini AI model",
                    "content": "Google's Gemini AI model represents a significant advancement in artificial intelligence technology.",
                },
                {
                    "id": "0",
                    "score": 0.92,
                    "title": "Gemini is a",
                    "content": "Gemini is a cutting edge large language model created by Google.",
                },
            ]
        }

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = mock_response_data
        mock_response.text = '{"records": [{"id": "3", "score": 0.95, "title": "Google\'s Gemini AI model", "content": "Google\'s Gemini AI model represents a significant advancement in artificial intelligence technology."}, {"id": "0", "score": 0.92, "title": "Gemini is a", "content": "Gemini is a cutting edge large language model created by Google."}]}'

        mock_logging = MagicMock()

        # Transform response
        from litellm.types.rerank import RerankResponse

        model_response = RerankResponse()

        result = self.config.transform_rerank_response(
            model=self.model,
            raw_response=mock_response,
            model_response=model_response,
            logging_obj=mock_logging,
        )

        # Verify response structure
        assert result.id == f"vertex_ai_rerank_{self.model}"
        assert len(result.results) == 2

        # Results should be sorted by relevance score (descending)
        assert result.results[0]["index"] == 3  # Highest score
        assert result.results[0]["relevance_score"] == 0.95
        assert result.results[1]["index"] == 0  # Second highest score
        assert result.results[1]["relevance_score"] == 0.92

        # Verify metadata
        assert result.meta["billed_units"]["search_units"] == 2

    def test_return_documents_false_flow(self):
        """Test rerank flow when return_documents=False (ID-only response)."""
        documents = ["doc1", "doc2", "doc3"]
        query = "test query"

        # Transform request with return_documents=False
        request_data = self.config.transform_rerank_request(
            model=self.model,
            optional_rerank_params={
                "query": query,
                "documents": documents,
                "return_documents": False,
            },
            headers={},
        )

        # Verify ignoreRecordDetailsInResponse is True
        assert request_data["ignoreRecordDetailsInResponse"] == True

        # Mock response with only IDs
        mock_response_data = {"records": [{"id": "1"}, {"id": "0"}, {"id": "2"}]}

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = mock_response_data
        mock_response.text = '{"records": [{"id": "1"}, {"id": "0"}, {"id": "2"}]}'

        mock_logging = MagicMock()

        # Transform response
        from litellm.types.rerank import RerankResponse

        model_response = RerankResponse()

        result = self.config.transform_rerank_response(
            model=self.model,
            raw_response=mock_response,
            model_response=model_response,
            logging_obj=mock_logging,
        )

        # Verify response structure with default scores
        assert len(result.results) == 3
        for result_item in result.results:
            assert (
                result_item["relevance_score"] == 1.0
            )  # Default score when details are ignored
            assert "index" in result_item

    def test_document_title_generation(self):
        """Test that document titles are generated correctly from content."""
        documents = [
            "This is a very long document with many words that should be truncated to only the first three words for the title",
            "Short doc",
            "Another document with multiple words here and more content",
        ]

        request_data = self.config.transform_rerank_request(
            model=self.model,
            optional_rerank_params={"query": "test query", "documents": documents},
            headers={},
        )

        # Verify title generation
        assert request_data["records"][0]["title"] == "This is a"  # First 3 words
        assert request_data["records"][1]["title"] == "Short doc"  # Less than 3 words
        assert (
            request_data["records"][2]["title"] == "Another document with"
        )  # First 3 words

    def test_dictionary_document_handling(self):
        """Test handling of dictionary-format documents."""
        documents = [
            {
                "text": "Gemini is a cutting edge large language model created by Google.",
                "title": "Custom Title 1",
            },
            {
                "text": "The Gemini zodiac symbol often depicts two figures standing side-by-side."
            },
            {
                "text": "Gemini is a constellation that can be seen in the night sky.",
                "title": "Custom Title 3",
            },
        ]

        request_data = self.config.transform_rerank_request(
            model=self.model,
            optional_rerank_params={"query": "test query", "documents": documents},
            headers={},
        )

        # Verify custom titles are used when provided
        assert request_data["records"][0]["title"] == "Custom Title 1"
        assert (
            request_data["records"][1]["title"] == "The Gemini zodiac"
        )  # Generated from first 3 words
        assert request_data["records"][2]["title"] == "Custom Title 3"

        # Verify content is extracted correctly
        assert (
            request_data["records"][0]["content"]
            == "Gemini is a cutting edge large language model created by Google."
        )
        assert (
            request_data["records"][1]["content"]
            == "The Gemini zodiac symbol often depicts two figures standing side-by-side."
        )
        assert (
            request_data["records"][2]["content"]
            == "Gemini is a constellation that can be seen in the night sky."
        )
