import json
from unittest.mock import MagicMock, Mock, patch

import pytest
from azure.core.exceptions import HttpResponseError

from mem0.configs.vector_stores.azure_ai_search import AzureAISearchConfig

# Import the AzureAISearch class and related models
from mem0.vector_stores.azure_ai_search import AzureAISearch


# Fixture to patch SearchClient and SearchIndexClient and create an instance of AzureAISearch.
@pytest.fixture
def mock_clients():
    with (
        patch("mem0.vector_stores.azure_ai_search.SearchClient") as MockSearchClient,
        patch("mem0.vector_stores.azure_ai_search.SearchIndexClient") as MockIndexClient,
        patch("mem0.vector_stores.azure_ai_search.AzureKeyCredential") as MockAzureKeyCredential,
    ):
        # Create mocked instances for search and index clients.
        mock_search_client = MockSearchClient.return_value
        mock_index_client = MockIndexClient.return_value

        # Mock the client._client._config.user_agent_policy.add_user_agent
        mock_search_client._client = MagicMock()
        mock_search_client._client._config.user_agent_policy.add_user_agent = Mock()
        mock_index_client._client = MagicMock()
        mock_index_client._client._config.user_agent_policy.add_user_agent = Mock()

        # Stub required methods on search_client.
        mock_search_client.upload_documents = Mock()
        mock_search_client.upload_documents.return_value = [{"status": True, "id": "doc1"}]
        mock_search_client.search = Mock()
        mock_search_client.delete_documents = Mock()
        mock_search_client.delete_documents.return_value = [{"status": True, "id": "doc1"}]
        mock_search_client.merge_or_upload_documents = Mock()
        mock_search_client.merge_or_upload_documents.return_value = [{"status": True, "id": "doc1"}]
        mock_search_client.get_document = Mock()
        mock_search_client.close = Mock()

        # Stub required methods on index_client.
        mock_index_client.create_or_update_index = Mock()
        mock_index_client.list_indexes = Mock()
        mock_index_client.list_index_names = Mock(return_value=[])
        mock_index_client.delete_index = Mock()
        # For col_info() we assume get_index returns an object with name and fields attributes.
        fake_index = Mock()
        fake_index.name = "test-index"
        fake_index.fields = ["id", "vector", "payload", "user_id", "run_id", "agent_id"]
        mock_index_client.get_index = Mock(return_value=fake_index)
        mock_index_client.close = Mock()

        yield mock_search_client, mock_index_client, MockAzureKeyCredential


@pytest.fixture
def azure_ai_search_instance(mock_clients):
    mock_search_client, mock_index_client, _ = mock_clients
    # Create an instance with dummy parameters.
    instance = AzureAISearch(
        service_name="test-service",
        collection_name="test-index",
        api_key="test-api-key",
        embedding_model_dims=3,
        compression_type="binary",  # testing binary quantization option
        use_float16=True,
    )
    # Return instance and clients for verification.
    return instance, mock_search_client, mock_index_client


# --- Tests for AzureAISearchConfig ---


def test_config_validation_valid():
    """Test valid configurations are accepted."""
    # Test minimal configuration
    config = AzureAISearchConfig(service_name="test-service", api_key="test-api-key", embedding_model_dims=768)
    assert config.collection_name == "mem0"  # Default value
    assert config.service_name == "test-service"
    assert config.api_key == "test-api-key"
    assert config.embedding_model_dims == 768
    assert config.compression_type is None
    assert config.use_float16 is False

    # Test with all optional parameters
    config = AzureAISearchConfig(
        collection_name="custom-index",
        service_name="test-service",
        api_key="test-api-key",
        embedding_model_dims=1536,
        compression_type="scalar",
        use_float16=True,
    )
    assert config.collection_name == "custom-index"
    assert config.compression_type == "scalar"
    assert config.use_float16 is True


def test_config_validation_invalid_compression_type():
    """Test that invalid compression types are rejected."""
    with pytest.raises(ValueError) as exc_info:
        AzureAISearchConfig(
            service_name="test-service",
            api_key="test-api-key",
            embedding_model_dims=768,
            compression_type="invalid-type",  # Not a valid option
        )
    assert "Invalid compression_type" in str(exc_info.value)


def test_config_validation_deprecated_use_compression():
    """Test that using the deprecated use_compression parameter raises an error."""
    with pytest.raises(ValueError) as exc_info:
        AzureAISearchConfig(
            service_name="test-service",
            api_key="test-api-key",
            embedding_model_dims=768,
            use_compression=True,  # Deprecated parameter
        )
    # Fix: Use a partial string match instead of exact match
    assert "use_compression" in str(exc_info.value)
    assert "no longer supported" in str(exc_info.value)


def test_config_validation_extra_fields():
    """Test that extra fields are rejected."""
    with pytest.raises(ValueError) as exc_info:
        AzureAISearchConfig(
            service_name="test-service",
            api_key="test-api-key",
            embedding_model_dims=768,
            unknown_parameter="value",  # Extra field
        )
    assert "Extra fields not allowed" in str(exc_info.value)
    assert "unknown_parameter" in str(exc_info.value)


# --- Tests for AzureAISearch initialization ---


def test_initialization(mock_clients):
    """Test AzureAISearch initialization with different parameters."""
    mock_search_client, mock_index_client, mock_azure_key_credential = mock_clients

    # Test with minimal parameters
    instance = AzureAISearch(
        service_name="test-service", collection_name="test-index", api_key="test-api-key", embedding_model_dims=768
    )

    # Verify initialization parameters
    assert instance.index_name == "test-index"
    assert instance.collection_name == "test-index"
    assert instance.embedding_model_dims == 768
    assert instance.compression_type == "none"  # Default when None is passed
    assert instance.use_float16 is False

    # Verify client creation
    mock_azure_key_credential.assert_called_with("test-api-key")
    assert "mem0" in mock_search_client._client._config.user_agent_policy.add_user_agent.call_args[0]
    assert "mem0" in mock_index_client._client._config.user_agent_policy.add_user_agent.call_args[0]

    # Verify index creation was called
    mock_index_client.create_or_update_index.assert_called_once()


def test_initialization_with_compression_types(mock_clients):
    """Test initialization with different compression types."""
    mock_search_client, mock_index_client, _ = mock_clients

    # Test with scalar compression
    instance = AzureAISearch(
        service_name="test-service",
        collection_name="scalar-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        compression_type="scalar",
    )
    assert instance.compression_type == "scalar"

    # Capture the index creation call
    args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
    index = args[0]
    # Verify scalar compression was configured
    assert hasattr(index.vector_search, "compressions")
    assert len(index.vector_search.compressions) > 0
    assert "ScalarQuantizationCompression" in str(type(index.vector_search.compressions[0]))

    # Test with binary compression
    instance = AzureAISearch(
        service_name="test-service",
        collection_name="binary-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        compression_type="binary",
    )
    assert instance.compression_type == "binary"

    # Capture the index creation call
    args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
    index = args[0]
    # Verify binary compression was configured
    assert hasattr(index.vector_search, "compressions")
    assert len(index.vector_search.compressions) > 0
    assert "BinaryQuantizationCompression" in str(type(index.vector_search.compressions[0]))

    # Test with no compression
    instance = AzureAISearch(
        service_name="test-service",
        collection_name="no-compression-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        compression_type=None,
    )
    assert instance.compression_type == "none"

    # Capture the index creation call
    args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
    index = args[0]
    # Verify no compression was configured
    assert hasattr(index.vector_search, "compressions")
    assert len(index.vector_search.compressions) == 0


def test_initialization_with_float_precision(mock_clients):
    """Test initialization with different float precision settings."""
    mock_search_client, mock_index_client, _ = mock_clients

    # Test with half precision (float16)
    instance = AzureAISearch(
        service_name="test-service",
        collection_name="float16-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        use_float16=True,
    )
    assert instance.use_float16 is True

    # Capture the index creation call
    args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
    index = args[0]
    # Find the vector field and check its type
    vector_field = next((f for f in index.fields if f.name == "vector"), None)
    assert vector_field is not None
    assert "Edm.Half" in vector_field.type

    # Test with full precision (float32)
    instance = AzureAISearch(
        service_name="test-service",
        collection_name="float32-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        use_float16=False,
    )
    assert instance.use_float16 is False

    # Capture the index creation call
    args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
    index = args[0]
    # Find the vector field and check its type
    vector_field = next((f for f in index.fields if f.name == "vector"), None)
    assert vector_field is not None
    assert "Edm.Single" in vector_field.type


# --- Tests for create_col method ---


def test_create_col(azure_ai_search_instance):
    """Test the create_col method creates an index with the correct configuration."""
    instance, _, mock_index_client = azure_ai_search_instance

    # create_col is called during initialization, so we check the call that was already made
    mock_index_client.create_or_update_index.assert_called_once()

    # Verify the index configuration
    args, _ = mock_index_client.create_or_update_index.call_args
    index = args[0]

    # Check basic properties
    assert index.name == "test-index"
    assert len(index.fields) == 6  # id, user_id, run_id, agent_id, vector, payload

    # Check that required fields are present
    field_names = [f.name for f in index.fields]
    assert "id" in field_names
    assert "vector" in field_names
    assert "payload" in field_names
    assert "user_id" in field_names
    assert "run_id" in field_names
    assert "agent_id" in field_names

    # Check that id is the key field
    id_field = next(f for f in index.fields if f.name == "id")
    assert id_field.key is True

    # Check vector search configuration
    assert index.vector_search is not None
    assert len(index.vector_search.profiles) == 1
    assert index.vector_search.profiles[0].name == "my-vector-config"
    assert index.vector_search.profiles[0].algorithm_configuration_name == "my-algorithms-config"

    # Check algorithms
    assert len(index.vector_search.algorithms) == 1
    assert index.vector_search.algorithms[0].name == "my-algorithms-config"
    assert "HnswAlgorithmConfiguration" in str(type(index.vector_search.algorithms[0]))

    # With binary compression and float16, we should have compression configuration
    assert len(index.vector_search.compressions) == 1
    assert index.vector_search.compressions[0].compression_name == "myCompression"
    assert "BinaryQuantizationCompression" in str(type(index.vector_search.compressions[0]))


def test_create_col_scalar_compression(mock_clients):
    """Test creating a collection with scalar compression."""
    mock_search_client, mock_index_client, _ = mock_clients

    AzureAISearch(
        service_name="test-service",
        collection_name="scalar-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        compression_type="scalar",
    )

    # Verify the index configuration
    args, _ = mock_index_client.create_or_update_index.call_args
    index = args[0]

    # Check compression configuration
    assert len(index.vector_search.compressions) == 1
    assert index.vector_search.compressions[0].compression_name == "myCompression"
    assert "ScalarQuantizationCompression" in str(type(index.vector_search.compressions[0]))

    # Check profile references compression
    assert index.vector_search.profiles[0].compression_name == "myCompression"


def test_create_col_no_compression(mock_clients):
    """Test creating a collection with no compression."""
    mock_search_client, mock_index_client, _ = mock_clients

    AzureAISearch(
        service_name="test-service",
        collection_name="no-compression-index",
        api_key="test-api-key",
        embedding_model_dims=768,
        compression_type=None,
    )

    # Verify the index configuration
    args, _ = mock_index_client.create_or_update_index.call_args
    index = args[0]

    # Check compression configuration - should be empty
    assert len(index.vector_search.compressions) == 0

    # Check profile doesn't reference compression
    assert index.vector_search.profiles[0].compression_name is None


# --- Tests for insert method ---


def test_insert_single(azure_ai_search_instance):
    """Test inserting a single vector."""
    instance, mock_search_client, _ = azure_ai_search_instance
    vectors = [[0.1, 0.2, 0.3]]
    payloads = [{"user_id": "user1", "run_id": "run1", "agent_id": "agent1"}]
    ids = ["doc1"]

    # Fix: Include status_code: 201 in mock response
    mock_search_client.upload_documents.return_value = [{"status": True, "id": "doc1", "status_code": 201}]

    instance.insert(vectors, payloads, ids)

    # Verify upload_documents was called correctly
    mock_search_client.upload_documents.assert_called_once()
    args, _ = mock_search_client.upload_documents.call_args
    documents = args[0]

    # Verify document structure
    assert len(documents) == 1
    assert documents[0]["id"] == "doc1"
    assert documents[0]["vector"] == [0.1, 0.2, 0.3]
    assert documents[0]["payload"] == json.dumps(payloads[0])
    assert documents[0]["user_id"] == "user1"
    assert documents[0]["run_id"] == "run1"
    assert documents[0]["agent_id"] == "agent1"


def test_insert_multiple(azure_ai_search_instance):
    """Test inserting multiple vectors in one call."""
    instance, mock_search_client, _ = azure_ai_search_instance

    # Create multiple vectors
    num_docs = 3
    vectors = [[float(i) / 10, float(i + 1) / 10, float(i + 2) / 10] for i in range(num_docs)]
    payloads = [{"user_id": f"user{i}", "content": f"Test content {i}"} for i in range(num_docs)]
    ids = [f"doc{i}" for i in range(num_docs)]

    # Configure mock to return success for all documents (fix: add status_code 201)
    mock_search_client.upload_documents.return_value = [
        {"status": True, "id": id_val, "status_code": 201} for id_val in ids
    ]

    # Insert the documents
    instance.insert(vectors, payloads, ids)

    # Verify upload_documents was called with correct documents
    mock_search_client.upload_documents.assert_called_once()
    args, _ = mock_search_client.upload_documents.call_args
    documents = args[0]

    # Verify all documents were included
    assert len(documents) == num_docs

    # Check first document
    assert documents[0]["id"] == "doc0"
    assert documents[0]["vector"] == [0.0, 0.1, 0.2]
    assert documents[0]["payload"] == json.dumps(payloads[0])
    assert documents[0]["user_id"] == "user0"

    # Check last document
    assert documents[2]["id"] == "doc2"
    assert documents[2]["vector"] == [0.2, 0.3, 0.4]
    assert documents[2]["payload"] == json.dumps(payloads[2])
    assert documents[2]["user_id"] == "user2"


def test_insert_with_error(azure_ai_search_instance):
    """Test insert when Azure returns an error for one or more documents."""
    instance, mock_search_client, _ = azure_ai_search_instance

    # Configure mock to return an error for one document
    mock_search_client.upload_documents.return_value = [{"status": False, "id": "doc1", "errorMessage": "Azure error"}]

    vectors = [[0.1, 0.2, 0.3]]
    payloads = [{"user_id": "user1"}]
    ids = ["doc1"]

    # Insert should raise an exception
    with pytest.raises(Exception) as exc_info:
        instance.insert(vectors, payloads, ids)

    assert "Insert failed for document doc1" in str(exc_info.value)

    # Configure mock to return mixed success/failure for multiple documents
    mock_search_client.upload_documents.return_value = [
        {"status": True, "id": "doc1"},  # This should not cause failure
        {"status": False, "id": "doc2", "errorMessage": "Azure error"},
    ]

    vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
    payloads = [{"user_id": "user1"}, {"user_id": "user2"}]
    ids = ["doc1", "doc2"]

    # Insert should raise an exception, but now check for doc2 failure
    with pytest.raises(Exception) as exc_info:
        instance.insert(vectors, payloads, ids)

    assert "Insert failed for document doc2" in str(exc_info.value) or "Insert failed for document doc1" in str(
        exc_info.value
    )


def test_insert_with_missing_payload_fields(azure_ai_search_instance):
    """Test inserting with payloads missing some of the expected fields."""
    instance, mock_search_client, _ = azure_ai_search_instance
    vectors = [[0.1, 0.2, 0.3]]
    payloads = [{"content": "Some content without user_id, run_id, or agent_id"}]
    ids = ["doc1"]

    # Mock successful response with a proper status_code
    mock_search_client.upload_documents.return_value = [
        {"id": "doc1", "status_code": 201}  # Simulating a successful response
    ]

    instance.insert(vectors, payloads, ids)

    # Verify upload_documents was called correctly
    mock_search_client.upload_documents.assert_called_once()
    args, _ = mock_search_client.upload_documents.call_args
    documents = args[0]
    # Verify document has payload but not the extra fields
    assert len(documents) == 1
    assert documents[0]["id"] == "doc1"
    assert documents[0]["vector"] == [0.1, 0.2, 0.3]
    assert documents[0]["payload"] == json.dumps(payloads[0])
    assert "user_id" not in documents[0]
    assert "run_id" not in documents[0]
    assert "agent_id" not in documents[0]


def test_insert_with_http_error(azure_ai_search_instance):
    """Test insert when Azure client throws an HTTP error."""
    instance, mock_search_client, _ = azure_ai_search_instance

    # Configure mock to raise an HttpResponseError
    mock_search_client.upload_documents.side_effect = HttpResponseError("Azure service error")

    vectors = [[0.1, 0.2, 0.3]]
    payloads = [{"user_id": "user1"}]
    ids = ["doc1"]

    # Insert should propagate the HTTP error
    with pytest.raises(HttpResponseError) as exc_info:
        instance.insert(vectors, payloads, ids)

    assert "Azure service error" in str(exc_info.value)


# --- Tests for search method ---


def test_search_basic(azure_ai_search_instance):
    """Test basic vector search without filters."""
    instance, mock_search_client, _ = azure_ai_search_instance

    # Ensure instance has a default vector_filter_mode
    instance.vector_filter_mode = "preFilter"

    # Configure mock to return search results
    mock_search_client.search.return_value = [
        {
            "id": "doc1",
            "@search.score": 0.95,
            "payload": json.dumps({"content": "Test content"}),
        }
    ]

    # Search with a vector
    query_text = "test query"  # Add a query string
    query_vector = [0.1, 0.2, 0.3]
    results = instance.search(query_text, query_vector, limit=5)  # Pass the query string

    # Verify search was called correctly
    mock_search_client.search.assert_called_once()
    _, kwargs = mock_search_client.search.call_args

    # Check parameters
    assert len(kwargs["vector_queries"]) == 1
    assert kwargs["vector_queries"][0].vector == query_vector
    assert kwargs["vector_queries"][0].k_nearest_neighbors == 5
    assert kwargs["vector_queries"][0].fields == "vector"
    assert kwargs["filter"] is None  # No filters
    assert kwargs["top"] == 5
    assert kwargs["vector_filter_mode"] == "preFilter"  # Now correctly set

    # Check results
    assert len(results) == 1
    assert results[0].id == "doc1"
    assert results[0].score == 0.95
    assert results[0].payload == {"content": "Test content"}
