"""
Tests for IBM watsonx.ai rerank transformation functionality.
"""

import json
import re
import uuid
from unittest.mock import MagicMock

import httpx
import pytest

from litellm.llms.watsonx.common_utils import (
    WatsonXAIError,
)
from litellm.llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig
from litellm.types.rerank import RerankResponse


class TestIBMWatsonXRerankTransform:
    def setup_method(self):
        self.config = IBMWatsonXRerankConfig()
        self.model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"

    def test_get_complete_url(self):
        """Test URL generation for IBM watsonx.ai rerank API."""

        api_base = "https://us-south.ml.cloud.ibm.com"
        model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
        url = self.config.get_complete_url(api_base, model)
        assert (
            url
            == "https://us-south.ml.cloud.ibm.com/ml/v1/text/rerank?version=2024-03-13"
        )

    def test_map_cohere_rerank_params_basic(self):
        """Test basic parameter mapping for IBM watsonx.ai rerank."""
        params = self.config.map_cohere_rerank_params(
            non_default_params={
                "query": "hello",
                "documents": ["hello", "world"],
                "top_n": 2,
                "return_documents": True,
                "max_tokens_per_doc": 100,
            },
            model="test",
            drop_params=False,
            query="hello",
            documents=["hello", "world"],
        )
        assert params["query"] == "hello"
        assert params["inputs"] == [{"text": "hello"}, {"text": "world"}]
        assert params["parameters"]["return_options"]["top_n"] == 2
        assert params["parameters"]["return_options"]["inputs"] is True
        assert params["parameters"]["truncate_input_tokens"] == 100

    def test_transform_rerank_request(self):
        """Test request transformation for IBM watsonx.ai format."""
        optional_params = {
            "query": "What is the capital of France?",
            "documents": [
                "Paris is the capital of France.",
                "France is a country in Europe.",
            ],
            "top_n": 2,
            "return_documents": True,
            "project_id": uuid.uuid4(),
        }

        request_body = self.config.transform_rerank_request(
            model="cross-encoder/ms-marco-minilm-l-12-v2",
            optional_rerank_params=optional_params,
            headers={},
        )

        assert request_body["model_id"] == "cross-encoder/ms-marco-minilm-l-12-v2"
        assert request_body["project_id"] is not None
        assert request_body["query"] == "What is the capital of France?"
        assert request_body["documents"] == optional_params["documents"]
        assert request_body["top_n"] == 2
        assert request_body["return_documents"] is True

    def test_transform_rerank_response_success(self):
        """Test successful response transformation."""
        # Mock IBM watsonx.ai response format
        response_data = {
            "model_id": self.model,
            "results": [
                {
                    "index": 0,
                    "score": 6.53515625,
                    "input": {
                        "text": "Python is great for beginners due to simple syntax."
                    },
                },
                {
                    "index": 1,
                    "score": -7.1875,
                    "input": {"text": "JavaScript runs in browsers and is versatile."},
                },
            ],
            "input_token_count": 62,
        }

        # Create mock httpx response
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = response_data
        mock_response.status_code = 200
        mock_response.headers = {}

        # Create mock logging object
        mock_logging = MagicMock()

        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
        # IBM watsonx.ai doesn't return "id", so it uses "model" as the id
        assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
        assert len(result.results) == 2
        assert result.results[0]["index"] == 0
        assert result.results[0]["relevance_score"] == 6.53515625
        assert (
            result.results[0]["document"]["text"]
            == "Python is great for beginners due to simple syntax."
        )
        assert result.results[1]["index"] == 1
        assert result.results[1]["relevance_score"] == -7.1875
        assert (
            result.results[1]["document"]["text"]
            == "JavaScript runs in browsers and is versatile."
        )

        # # Verify metadata
        assert result.meta["tokens"]["input_tokens"] == 62

    def test_transform_rerank_response_without_documents(self):
        """Test response transformation when return_documents is False."""
        response_data = {
            "model_id": self.model,
            "results": [
                {
                    "index": 0,
                    "score": 6.53515625,
                },
                {
                    "index": 1,
                    "score": -7.1875,
                },
            ],
            "input_token_count": 62,
        }

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = response_data
        mock_response.status_code = 200
        mock_response.headers = {}

        mock_logging = MagicMock()
        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
        # IBM watsonx.ai doesn't return "id", so it uses "model" as the id
        assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
        assert len(result.results) == 2

        assert result.results[0]["index"] == 0
        assert result.results[0]["relevance_score"] == 6.53515625
        assert "document" not in result.results[0]

        assert result.results[1]["index"] == 1
        assert result.results[1]["relevance_score"] == -7.1875
        assert "document" not in result.results[1]

    def test_transform_rerank_response_missing_results(self):
        """Test that missing results raises ValueError."""
        response_data = {
            "model": self.model,
            "usage": {"total_tokens": 10},
        }

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.return_value = response_data
        mock_response.status_code = 200
        mock_response.headers = {}

        mock_logging = MagicMock()
        model_response = RerankResponse()

        expected_error_msg = re.escape("No results found")

        with pytest.raises(ValueError, match=expected_error_msg):
            self.config.transform_rerank_response(
                model=self.model,
                raw_response=mock_response,
                model_response=model_response,
                logging_obj=mock_logging,
            )

    def test_transform_rerank_response_invalid_json(self):
        """Test error handling for invalid JSON response."""
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0)
        mock_response.text = "Invalid JSON response"
        mock_response.status_code = 500
        mock_response.headers = {}

        mock_logging = MagicMock()
        model_response = RerankResponse()

        expected_error_msg = re.escape("Failed to parse response")

        with pytest.raises(Exception, match=expected_error_msg):
            self.config.transform_rerank_response(
                model=self.model,
                raw_response=mock_response,
                model_response=model_response,
                logging_obj=mock_logging,
            )

    def test_get_supported_cohere_rerank_params(self):
        """Test getting supported parameters for IBM watsonx.ai rerank."""
        supported_params = self.config.get_supported_cohere_rerank_params(self.model)
        assert "query" in supported_params
        assert "documents" in supported_params
        assert "top_n" in supported_params
        assert "return_documents" in supported_params
        assert "max_tokens_per_doc" in supported_params
        assert len(supported_params) == 5
