# 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 NeptuneServer class."""

import pytest
from awslabs.amazon_neptune_mcp_server.models import GraphSchema
from awslabs.amazon_neptune_mcp_server.neptune import NeptuneServer
from unittest.mock import MagicMock, patch


@pytest.mark.asyncio
class TestNeptuneServer:
    """Test class for the NeptuneServer functionality."""

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_init_neptune_db(self, mock_neptune_db):
        """Test initialization of NeptuneServer with a Neptune Database endpoint.
        This test verifies that:
        1. The Neptune Database endpoint is correctly parsed
        2. NeptuneDatabase is initialized with the correct parameters
        3. The graph attribute is set to the NeptuneDatabase instance.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_neptune_db.return_value = mock_db_instance
        endpoint = 'neptune-db://test-endpoint'

        # Act
        server = NeptuneServer(endpoint, use_https=True, port=8182)

        # Assert
        assert server.graph == mock_db_instance
        mock_neptune_db.assert_called_once_with('test-endpoint', 8182, use_https=True)

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneAnalytics')
    async def test_init_neptune_analytics(self, mock_neptune_analytics):
        """Test initialization of NeptuneServer with a Neptune Analytics endpoint.
        This test verifies that:
        1. The Neptune Analytics endpoint is correctly parsed
        2. NeptuneAnalytics is initialized with the correct graph ID
        3. The graph attribute is set to the NeptuneAnalytics instance.
        """
        # Arrange
        mock_analytics_instance = MagicMock()
        mock_neptune_analytics.return_value = mock_analytics_instance
        endpoint = 'neptune-graph://g-1234567890'

        # Act
        server = NeptuneServer(endpoint)

        # Assert
        assert server.graph == mock_analytics_instance
        mock_neptune_analytics.assert_called_once_with('g-1234567890', endpoint_url=None)

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneAnalytics')
    async def test_init_neptune_analytics_custom_endpoint(self, mock_neptune_analytics):
        """Test initialization with a custom endpoint.

        Any non-graphId value is treated as a custom endpoint with a placeholder graph ID:
        neptune-graph://localhost:9100 → graphId='g-1234567890', endpoint_url='http://localhost:9100'.
        """
        mock_analytics_instance = MagicMock()
        mock_neptune_analytics.return_value = mock_analytics_instance
        endpoint = 'neptune-graph://localhost:9100'

        server = NeptuneServer(endpoint)

        assert server.graph == mock_analytics_instance
        mock_neptune_analytics.assert_called_once_with(
            'g-1234567890', endpoint_url='http://localhost:9100'
        )

    async def test_init_invalid_endpoint_format(self):
        """Test that NeptuneServer initialization fails with an invalid endpoint format.
        This test verifies that:
        1. When an endpoint with an invalid format is provided, a ValueError is raised
        2. The error message correctly indicates the expected format.
        """
        # Act & Assert
        with pytest.raises(
            ValueError,
            match='You must provide an endpoint to create a NeptuneServer as either neptune-db',
        ):
            NeptuneServer('invalid-endpoint')

    async def test_init_empty_endpoint(self):
        """Test that NeptuneServer initialization fails with an empty endpoint.
        This test verifies that:
        1. When an empty endpoint is provided, a ValueError is raised
        2. The error message correctly indicates that an endpoint is required.
        """
        # Act & Assert
        with pytest.raises(
            ValueError, match='You must provide an endpoint to create a NeptuneServer'
        ):
            NeptuneServer('')

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_status_available(self, mock_neptune_db):
        """Test that status() returns "Available" when the database is available.
        This test verifies that:
        1. A test query is executed to check database availability
        2. When the query succeeds, "Available" is returned.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_db_instance.query_opencypher.return_value = {'result': 1}
        mock_neptune_db.return_value = mock_db_instance

        server = NeptuneServer('neptune-db://test-endpoint')

        # Act
        status = server.status()

        # Assert
        assert status == 'Available'
        mock_db_instance.query_opencypher.assert_called_once_with('RETURN 1', None)

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_status_unavailable(self, mock_neptune_db):
        """Test that status() returns "Unavailable" when the database is unavailable.
        This test verifies that:
        1. A test query is executed to check database availability
        2. When the query fails, "Unavailable" is returned
        3. The exception is properly handled.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_db_instance.query_opencypher.side_effect = Exception('Connection error')
        mock_neptune_db.return_value = mock_db_instance

        server = NeptuneServer('neptune-db://test-endpoint')

        # Act
        status = server.status()

        # Assert
        assert status == 'Unavailable'
        mock_db_instance.query_opencypher.assert_called_once_with('RETURN 1', None)

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_schema(self, mock_neptune_db):
        """Test that schema() correctly returns the graph schema.
        This test verifies that:
        1. The get_schema method is called on the graph instance
        2. The result from the graph's get_schema method is returned unchanged.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_schema = GraphSchema(nodes=[], relationships=[], relationship_patterns=[])
        mock_db_instance.get_schema.return_value = mock_schema
        mock_neptune_db.return_value = mock_db_instance

        server = NeptuneServer('neptune-db://test-endpoint')

        # Act
        schema = server.schema()

        # Assert
        assert schema == mock_schema
        mock_db_instance.get_schema.assert_called_once()

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_query_opencypher(self, mock_neptune_db):
        """Test that query_opencypher correctly executes an openCypher query without parameters.
        This test verifies that:
        1. The query_opencypher method is called on the graph instance with the correct query
        2. The result from the graph's query_opencypher method is returned unchanged.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_result = {'results': [{'n': {'id': '1'}}]}
        mock_db_instance.query_opencypher.return_value = mock_result
        mock_neptune_db.return_value = mock_db_instance

        server = NeptuneServer('neptune-db://test-endpoint')

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

        # Assert
        assert result == mock_result
        mock_db_instance.query_opencypher.assert_called_once_with(
            'MATCH (n) RETURN n LIMIT 1', None
        )

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_query_opencypher_with_parameters(self, mock_neptune_db):
        """Test that query_opencypher correctly executes an openCypher query with parameters.
        This test verifies that:
        1. The query_opencypher method is called on the graph instance with the correct query and parameters
        2. The result from the graph's query_opencypher method is returned unchanged.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_result = {'results': [{'n': {'id': '1'}}]}
        mock_db_instance.query_opencypher.return_value = mock_result
        mock_neptune_db.return_value = mock_db_instance

        server = NeptuneServer('neptune-db://test-endpoint')
        parameters = {'id': '1'}

        # Act
        result = server.query_opencypher('MATCH (n) WHERE n.id = $id RETURN n', parameters)

        # Assert
        assert result == mock_result
        mock_db_instance.query_opencypher.assert_called_once_with(
            'MATCH (n) WHERE n.id = $id RETURN n', parameters
        )

    @patch('awslabs.amazon_neptune_mcp_server.neptune.NeptuneDatabase')
    async def test_query_gremlin(self, mock_neptune_db):
        """Test that query_gremlin correctly executes a Gremlin query.

        This test verifies that:
        1. The query_gremlin method is called on the graph instance with the correct query
        2. The result from the graph's query_gremlin method is returned unchanged.
        """
        # Arrange
        mock_db_instance = MagicMock()
        mock_result = {'results': [{'id': '1'}]}
        mock_db_instance.query_gremlin.return_value = mock_result
        mock_neptune_db.return_value = mock_db_instance

        server = NeptuneServer('neptune-db://test-endpoint')

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

        # Assert
        assert result == mock_result
        mock_db_instance.query_gremlin.assert_called_once_with('g.V().limit(1)')
