# 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 NeptuneAnalytics 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.analytics import NeptuneAnalytics
from awslabs.amazon_neptune_mcp_server.models import (
    GraphSchema,
)
from unittest.mock import MagicMock, Mock, patch


@pytest.mark.asyncio
class TestNeptuneAnalytics:
    """Test class for the NeptuneAnalytics functionality."""

    @patch('boto3.Session')
    async def test_init_success(self, mock_session):
        """Test successful initialization of NeptuneAnalytics.

        This test verifies that:
        1. The boto3 Session is created correctly
        2. The client is created with the correct service name
        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(
            NeptuneAnalytics,
            '_refresh_schema',
            return_value=GraphSchema(nodes=[], relationships=[], relationship_patterns=[]),
        ):
            # Act
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

            # Assert
            mock_session.assert_called_once()
            mock_session_instance.client.assert_called_once_with(
                'neptune-graph', config=USER_AGENT_CONFIG
            )
            assert analytics.client == mock_client
            assert analytics.graph_identifier == 'test-graph-id'

    @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 service name.
        """
        # 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(
            NeptuneAnalytics,
            '_refresh_schema',
            return_value=GraphSchema(nodes=[], relationships=[], relationship_patterns=[]),
        ):
            # Act
            NeptuneAnalytics(
                graph_identifier='test-graph-id', credentials_profile_name='test-profile'
            )

            # Assert
            mock_session.assert_called_once_with(profile_name='test-profile')
            mock_session_instance.client.assert_called_once_with(
                'neptune-graph', config=USER_AGENT_CONFIG
            )

    @patch('boto3.Session')
    async def test_init_with_endpoint_url(self, mock_session):
        """Test initialization with a custom endpoint URL.

        Verifies that endpoint_url is passed to the client and inject_host_prefix is disabled.
        """
        mock_session_instance = MagicMock()
        mock_client = MagicMock()
        mock_session_instance.client.return_value = mock_client
        mock_session.return_value = mock_session_instance

        with patch.object(
            NeptuneAnalytics,
            '_refresh_schema',
            return_value=GraphSchema(nodes=[], relationships=[], relationship_patterns=[]),
        ):
            analytics = NeptuneAnalytics(
                graph_identifier='g-1234567890',
                endpoint_url='http://localhost:9100',
            )

            mock_session_instance.client.assert_called_once()
            call_kwargs = mock_session_instance.client.call_args
            assert call_kwargs[1]['endpoint_url'] == 'http://localhost:9100'
            assert call_kwargs[1]['config'].inject_host_prefix is False
            assert analytics.client == mock_client

    @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'
        ):
            NeptuneAnalytics(graph_identifier='test-graph-id')

    @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(
            NeptuneAnalytics, '_refresh_schema', side_effect=Exception('Schema refresh error')
        ):
            # Act & Assert
            with pytest.raises(NeptuneException) as exc_info:
                NeptuneAnalytics(graph_identifier='test-graph-id')

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

    @patch('boto3.Session')
    async def test_refresh_schema(self, mock_session):
        """Test schema refresh functionality.
        This test verifies that:
        1. The query_opencypher method is called with the pg_schema query
        2. The schema data is correctly processed from the response
        3. The schema is stored in the instance and 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 response
        mock_schema_data = {
            'labelTriples': [
                {'~from': 'Person', '~type': 'KNOWS', '~to': 'Person'},
                {'~from': 'Person', '~type': 'ACTED_IN', '~to': 'Movie'},
            ],
            'nodeLabels': ['Person', 'Movie'],
            'edgeLabels': ['KNOWS', 'ACTED_IN'],
            'nodeLabelDetails': {
                'Person': {
                    'properties': {
                        'name': {'datatypes': ['STRING']},
                        'age': {'datatypes': ['INTEGER']},
                    }
                },
                'Movie': {
                    'properties': {
                        'title': {'datatypes': ['STRING']},
                        'year': {'datatypes': ['INTEGER']},
                    }
                },
            },
            'edgeLabelDetails': {
                'KNOWS': {'properties': {'since': {'datatypes': ['DATE']}}},
                'ACTED_IN': {'properties': {'role': {'datatypes': ['STRING']}}},
            },
        }

        # Create the analytics instance
        with patch.object(NeptuneAnalytics, '_refresh_schema'):
            NeptuneAnalytics(graph_identifier='test-graph-id')

        # Create a new instance for testing the actual method
        analytics2 = NeptuneAnalytics.__new__(NeptuneAnalytics)
        analytics2.graph_identifier = 'test-graph-id'
        analytics2.schema = None

        # Mock query_opencypher to return the schema data
        analytics2.query_opencypher = MagicMock(return_value=[{'schema': mock_schema_data}])

        # Act
        schema = analytics2._refresh_schema()

        # Assert
        analytics2.query_opencypher.assert_called_once()
        assert len(schema.nodes) == 2
        assert len(schema.relationships) == 2
        assert len(schema.relationship_patterns) == 2

        # Check that the schema was stored in the instance
        assert analytics2.schema == schema

        # Verify node properties
        person_node = next((n for n in schema.nodes if n.labels == 'Person'), None)
        assert person_node is not None
        assert len(person_node.properties) == 2
        assert any(p.name == 'name' and p.type == ['STRING'] for p in person_node.properties)
        assert any(p.name == 'age' and p.type == ['INTEGER'] for p in person_node.properties)

        # Verify relationship properties
        knows_rel = next((r for r in schema.relationships if r.type == 'KNOWS'), None)
        assert knows_rel is not None
        assert len(knows_rel.properties) == 1
        assert knows_rel.properties[0].name == 'since'
        assert knows_rel.properties[0].type == ['DATE']

        # Verify relationship patterns
        knows_pattern = next(
            (p for p in schema.relationship_patterns if p.relation == 'KNOWS'), None
        )
        assert knows_pattern is not None
        assert knows_pattern.left_node == 'Person'
        assert knows_pattern.right_node == 'Person'

    @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(NeptuneAnalytics, '_refresh_schema', return_value=mock_schema):
            # Create the analytics instance
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

            # Act
            result = analytics.get_schema()

            # Assert - just verify the result is the 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(NeptuneAnalytics, '_refresh_schema', return_value=mock_schema):
            # Create the analytics instance
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

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

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

            # Act
            result = analytics.get_schema()

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

    @patch('boto3.Session')
    async def test_query_opencypher_success(self, mock_session):
        """Test successful execution of openCypher queries.
        This test verifies that:
        1. The execute_query API is called with the correct parameters
        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_payload = Mock()
        mock_payload.read.return_value = json.dumps({'results': [{'n': {'id': '1'}}]}).encode(
            'utf-8'
        )
        mock_client.execute_query.return_value = {'payload': mock_payload}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneAnalytics, '_refresh_schema'):
            # Create the analytics instance
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

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

            # Assert
            mock_client.execute_query.assert_called_once_with(
                graphIdentifier='test-graph-id',
                queryString=query,
                parameters={},
                language='OPEN_CYPHER',
            )
            assert result == [{'n': {'id': '1'}}]

    @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_query API is called with the correct parameters
        2. The parameters are passed correctly
        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_payload = Mock()
        mock_payload.read.return_value = json.dumps({'results': [{'n': {'id': '1'}}]}).encode(
            'utf-8'
        )
        mock_client.execute_query.return_value = {'payload': mock_payload}

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneAnalytics, '_refresh_schema'):
            # Create the analytics instance
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

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

            # Assert
            mock_client.execute_query.assert_called_once_with(
                graphIdentifier='test-graph-id',
                queryString=query,
                parameters=params,
                language='OPEN_CYPHER',
            )
            assert result == [{'n': {'id': '1'}}]

    @patch('boto3.Session')
    async def test_query_opencypher_error(self, mock_session):
        """Test handling of errors in openCypher queries.
        This test verifies that:
        1. API errors 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 the API to raise an exception
        mock_client.execute_query.side_effect = Exception('Query error')

        # Mock _refresh_schema to avoid actual API calls during init
        with patch.object(NeptuneAnalytics, '_refresh_schema'):
            # Create the analytics instance
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

            # Act & Assert
            with pytest.raises(NeptuneException) as exc_info:
                analytics.query_opencypher('MATCH (n) RETURN n')

            # Check the exception details
            assert 'An error occurred while executing the query' in exc_info.value.message
            assert 'Query error' in exc_info.value.details

    @patch('boto3.Session')
    async def test_query_gremlin_not_supported(self, mock_session):
        """Test that Gremlin queries are not supported.
        This test verifies that:
        1. Calling query_gremlin raises NotImplementedError
        2. The error message indicates that Gremlin is not supported.
        """
        # 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(NeptuneAnalytics, '_refresh_schema'):
            # Create the analytics instance
            analytics = NeptuneAnalytics(graph_identifier='test-graph-id')

            # Act & Assert
            with pytest.raises(
                NotImplementedError,
                match='Gremlin queries are not supported for Neptune Analytics graphs',
            ):
                analytics.query_gremlin('g.V().limit(1)')
