# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the NeptuneDatabase class."""

import json
import pytest
from awslabs.amazon_neptune_mcp_server.constants import USER_AGENT_CONFIG
from awslabs.amazon_neptune_mcp_server.exceptions import NeptuneException
from awslabs.amazon_neptune_mcp_server.graph_store.database import NeptuneDatabase
from awslabs.amazon_neptune_mcp_server.models import GraphSchema
from unittest.mock import MagicMock, patch


@pytest.mark.asyncio
class TestNeptuneDatabase:
    """Test class for the NeptuneDatabase functionality."""

    @patch('boto3.Session')
    async def test_init_success(self, mock_session):
        """Test successful initialization of NeptuneDatabase.
        This test verifies that:
        1. The boto3 Session is created correctly
        2. The client is created with the correct parameters
        3. The schema is refreshed during initialization.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock _refresh_schema to avoid actual API calls
        with patch.object(
            NeptuneDatabase,
            '_refresh_schema',
            return_value=GraphSchema(nodes=[], relationships=[], relationship_patterns=[]),
        ):
            # Act
            db = NeptuneDatabase(host='test-endpoint', port=8182, use_https=True)

            # Assert
            mock_session.assert_called_once()
            mock_session_instance.client.assert_called_once_with(
                'neptunedata', config=USER_AGENT_CONFIG, endpoint_url='https://test-endpoint:8182'
            )
            assert db.client == mock_client

    @patch('boto3.Session')
    async def test_init_with_credentials_profile(self, mock_session):
        """Test initialization with a credentials profile.
        This test verifies that:
        1. The boto3 Session is created with the specified profile name
        2. The client is created with the correct parameters.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock _refresh_schema to avoid actual API calls
        with patch.object(
            NeptuneDatabase,
            '_refresh_schema',
            return_value=GraphSchema(nodes=[], relationships=[], relationship_patterns=[]),
        ):
            # Act
            NeptuneDatabase(
                host='test-endpoint',
                port=8182,
                use_https=True,
                credentials_profile_name='test-profile',
            )

            # Assert
            mock_session.assert_called_once_with(profile_name='test-profile')
            mock_session_instance.client.assert_called_once_with(
                'neptunedata', config=USER_AGENT_CONFIG, endpoint_url='https://test-endpoint:8182'
            )

    @patch('boto3.Session')
    async def test_init_with_http(self, mock_session):
        """Test initialization with HTTP instead of HTTPS.
        This test verifies that:
        1. The client is created with an HTTP endpoint URL when use_https is False.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock _refresh_schema to avoid actual API calls
        with patch.object(
            NeptuneDatabase,
            '_refresh_schema',
            return_value=GraphSchema(nodes=[], relationships=[], relationship_patterns=[]),
        ):
            # Act
            NeptuneDatabase(host='test-endpoint', port=8182, use_https=False)

            # Assert
            mock_session_instance.client.assert_called_once_with(
                'neptunedata', config=USER_AGENT_CONFIG, endpoint_url='http://test-endpoint:8182'
            )

    @patch('boto3.Session')
    async def test_init_session_error(self, mock_session):
        """Test handling of session creation errors.
        This test verifies that:
        1. Errors during session creation are properly caught and re-raised
        2. The error message is appropriate.
        """
        # Arrange
        mock_session.side_effect = Exception('Auth error')

        # Act & Assert
        with pytest.raises(
            ValueError, match='Could not load credentials to authenticate with AWS client'
        ):
            NeptuneDatabase(host='test-endpoint')

    @patch('boto3.Session')
    async def test_init_refresh_schema_error(self, mock_session):
        """Test handling of schema refresh errors.

        This test verifies that:
        1. Errors during schema refresh are properly caught and re-raised as NeptuneException
        2. The error message is appropriate
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock _refresh_schema to raise an exception
        with patch.object(
            NeptuneDatabase, '_refresh_schema', side_effect=Exception('Schema refresh error')
        ):
            # Act & Assert
            with pytest.raises(NeptuneException) as exc_info:
                NeptuneDatabase(host='test-endpoint')

            # Check the exception details
            assert 'Could not get schema for Neptune database' in exc_info.value.message

    @patch('boto3.Session')
    async def test_get_summary_success(self, mock_session):
        """Test successful retrieval of graph summary.
        This test verifies that:
        1. The get_propertygraph_summary API is called
        2. The summary data is correctly extracted from the response.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API response
        mock_summary = {'nodeLabels': ['Person', 'Movie'], 'edgeLabels': ['ACTED_IN', 'DIRECTED']}
        mock_client.get_propertygraph_summary.return_value = {
            'payload': {'graphSummary': mock_summary}
        }

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            result = db._get_summary()

            # Assert
            mock_client.get_propertygraph_summary.assert_called_once()
            assert result == mock_summary

    @patch('boto3.Session')
    async def test_get_summary_api_error(self, mock_session):
        """Test handling of API errors in get_summary.
        This test verifies that:
        1. API errors are properly caught and re-raised as NeptuneException
        2. The error message indicates the Summary API is not available.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API to raise an exception
        mock_client.get_propertygraph_summary.side_effect = Exception('API error')

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act & Assert
            with pytest.raises(NeptuneException) as exc_info:
                db._get_summary()

            # Check the exception details
            assert 'Summary API is not available' in exc_info.value.message
            assert 'API error' in exc_info.value.details

    @patch('boto3.Session')
    async def test_get_summary_invalid_response(self, mock_session):
        """Test handling of invalid responses in get_summary.
        This test verifies that:
        1. Invalid responses are properly caught and re-raised as NeptuneException
        2. The error message indicates the response was invalid.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API to return an invalid response
        class MockResponse:
            def __init__(self):
                self.payload = {}  # Missing graphSummary
                self.content = b'Invalid response'

            def __getitem__(self, key):
                return getattr(self, key)

        mock_client.get_propertygraph_summary.return_value = MockResponse()

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act & Assert
            with pytest.raises(NeptuneException) as exc_info:
                db._get_summary()

            # Check the exception details
            assert 'Summary API did not return a valid response' in exc_info.value.message

    @patch('boto3.Session')
    async def test_get_labels(self, mock_session):
        """Test retrieval of node and edge labels.
        This test verifies that:
        1. The _get_summary method is called
        2. Node and edge labels are correctly extracted from the summary.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Mock _get_summary
            mock_summary = {
                'nodeLabels': ['Person', 'Movie'],
                'edgeLabels': ['ACTED_IN', 'DIRECTED'],
            }
            with patch.object(db, '_get_summary', return_value=mock_summary):
                # Act
                n_labels, e_labels = db._get_labels()

                # Assert
                assert n_labels == ['Person', 'Movie']
                assert e_labels == ['ACTED_IN', 'DIRECTED']

    @patch('boto3.Session')
    async def test_query_opencypher_without_params(self, mock_session):
        """Test execution of openCypher queries without parameters.
        This test verifies that:
        1. The execute_open_cypher_query API is called with the correct query
        2. The result is correctly extracted from the response.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API response
        mock_result = [{'n': {'id': '1'}}]
        mock_client.execute_open_cypher_query.return_value = {'result': mock_result}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            result = db.query_opencypher('MATCH (n) RETURN n LIMIT 1')

            # Assert
            mock_client.execute_open_cypher_query.assert_called_once_with(
                openCypherQuery='MATCH (n) RETURN n LIMIT 1'
            )
            assert result == mock_result

    @patch('boto3.Session')
    async def test_query_opencypher_with_params(self, mock_session):
        """Test execution of openCypher queries with parameters.
        This test verifies that:
        1. The execute_open_cypher_query API is called with the correct query and parameters
        2. The parameters are properly JSON-encoded
        3. The result is correctly extracted from the response.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API response
        mock_result = [{'n': {'id': '1'}}]
        mock_client.execute_open_cypher_query.return_value = {'result': mock_result}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            params = {'id': '1'}
            result = db.query_opencypher('MATCH (n) WHERE n.id = $id RETURN n', params)

            # Assert
            mock_client.execute_open_cypher_query.assert_called_once_with(
                openCypherQuery='MATCH (n) WHERE n.id = $id RETURN n',
                parameters=json.dumps(params),
            )
            assert result == mock_result

    @patch('boto3.Session')
    async def test_query_opencypher_results_format(self, mock_session):
        """Test handling of different result formats in openCypher queries.
        This test verifies that:
        1. The method correctly handles responses with 'results' instead of 'result'.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API response with 'results' instead of 'result'
        mock_results = [{'n': {'id': '1'}}]
        mock_client.execute_open_cypher_query.return_value = {'results': mock_results}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            result = db.query_opencypher('MATCH (n) RETURN n LIMIT 1')

            # Assert
            assert result == mock_results

    @patch('boto3.Session')
    async def test_query_gremlin(self, mock_session):
        """Test execution of Gremlin queries.
        This test verifies that:
        1. The execute_gremlin_query API is called with the correct query
        2. The serializer parameter is correctly set
        3. The result is correctly extracted from the response.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API response
        mock_result = [{'id': '1'}]
        mock_client.execute_gremlin_query.return_value = {'result': mock_result}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            result = db.query_gremlin('g.V().limit(1)')

            # Assert
            mock_client.execute_gremlin_query.assert_called_once()
            assert result == mock_result

    @patch('boto3.Session')
    async def test_query_gremlin_results_format(self, mock_session):
        """Test handling of different result formats in Gremlin queries.
        This test verifies that:
        1. The method correctly handles responses with 'results' instead of 'result'.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Mock the API response with 'results' instead of 'result'
        mock_results = [{'id': '1'}]
        mock_client.execute_gremlin_query.return_value = {'results': mock_results}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema'):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            result = db.query_gremlin('g.V().limit(1)')

            # Assert
            assert result == mock_results

    @patch('boto3.Session')
    async def test_get_schema_cached(self, mock_session):
        """Test that get_schema returns cached schema when available.
        This test verifies that:
        1. When schema is already cached, _refresh_schema is not called
        2. The cached schema is returned.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Create a mock schema
        mock_schema = GraphSchema(nodes=[], relationships=[], relationship_patterns=[])

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema', return_value=mock_schema):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Act
            result = db.get_schema()

            # Assert - just verify the result is the mock schema
            assert result == mock_schema
            assert result == mock_schema

    @patch('boto3.Session')
    async def test_get_schema_refresh(self, mock_session):
        """Test that get_schema refreshes schema when not cached.
        This test verifies that:
        1. When schema is not cached, _refresh_schema is called
        2. The refreshed schema is returned.
        """
        # Arrange
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        # Create a mock schema
        mock_schema = GraphSchema(nodes=[], relationships=[], relationship_patterns=[])

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneDatabase, '_refresh_schema', return_value=mock_schema):
            # Create the database instance
            db = NeptuneDatabase(host='test-endpoint')

            # Set schema to None to force refresh
            db.schema = None

            # Reset the mock to verify it's called again
            NeptuneDatabase._refresh_schema.reset_mock()
            NeptuneDatabase._refresh_schema.return_value = mock_schema

            # Act
            result = db.get_schema()

            # Assert
            NeptuneDatabase._refresh_schema.assert_called_once()
            assert result == mock_schema
