"""Tests for batch query tool."""

from unittest.mock import AsyncMock, patch

import pytest

from open_targets_platform_mcp.model.result import BatchQueryResult, QueryResult, QueryResultStatus
from open_targets_platform_mcp.tools.batch_query.batch_query import _batch_query_impl

# Use the internal implementation function directly for testing
batch_query_fn = _batch_query_impl


class TestBatchQueryOpenTargetsGraphQL:
    """Tests for batch_query_open_targets_graphql function."""

    @pytest.mark.asyncio
    async def test_batch_query_success(self, batch_query_string, batch_variables_with_key):
        """Test successful batch query execution."""
        # Mock execute_graphql_query to return success for each call
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.side_effect = [
                QueryResult.create_success({"target": {"id": "ENSG00000141510", "approvedSymbol": "TP53"}}),
                QueryResult.create_success({"target": {"id": "ENSG00000012048", "approvedSymbol": "BRCA1"}}),
                QueryResult.create_success({"target": {"id": "ENSG00000139618", "approvedSymbol": "BRCA2"}}),
            ]

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
            )

        assert isinstance(result, BatchQueryResult)
        assert result.summary.total == 3
        assert result.summary.successful == 3
        assert result.summary.failed == 0

        # Check that results are in the list with correct keys
        result_dict = {r.key: r for r in result.results if r.key is not None}
        assert "ENSG00000141510" in result_dict
        assert "ENSG00000012048" in result_dict
        assert "ENSG00000139618" in result_dict

    @pytest.mark.asyncio
    async def test_batch_query_empty_variables_list(self, batch_query_string):
        """Test that empty variables_list returns error."""
        result = await batch_query_fn(query_string=batch_query_string, variables_list=[], key_field="ensemblId")

        assert isinstance(result, QueryResult)
        assert result.status == QueryResultStatus.ERROR
        assert "cannot be empty" in str(result.message)

    @pytest.mark.asyncio
    async def test_batch_query_missing_key_field(self, batch_query_string):
        """Test handling when key_field is missing from variables."""
        variables_list = [
            {"ensemblId": "ENSG00000141510"},
            {"wrongField": "value"},  # Missing ensemblId
            {"ensemblId": "ENSG00000139618"},
        ]

        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.side_effect = [
                QueryResult.create_success({"target": {"id": "ENSG00000141510"}}),
                QueryResult.create_success({"target": {"id": "ENSG00000139618"}}),
            ]

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=variables_list,
                key_field="ensemblId",
            )

        assert isinstance(result, BatchQueryResult)
        assert result.summary.total == 3
        assert result.summary.successful == 2
        assert result.summary.failed == 1

        # Check that missing key_field entry exists and has error
        error_result = result.results[1]  # Index 1 is the one with missing key
        assert error_result.key is None
        assert error_result.result.status == QueryResultStatus.ERROR
        assert "not found" in str(error_result.result.message)

    @pytest.mark.asyncio
    async def test_batch_query_partial_failures(self, batch_query_string, batch_variables_with_key):
        """Test batch query with some queries failing."""
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.side_effect = [
                QueryResult.create_success({"target": {"id": "ENSG00000141510"}}),
                QueryResult.create_error("Query failed"),
                QueryResult.create_success({"target": {"id": "ENSG00000139618"}}),
            ]

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
            )

        assert isinstance(result, BatchQueryResult)
        assert result.summary.total == 3
        assert result.summary.successful == 2
        assert result.summary.failed == 1

        # Check the failed query result
        result_dict = {r.key: r for r in result.results if r.key is not None}
        assert result_dict["ENSG00000012048"].result.status == QueryResultStatus.ERROR

    @pytest.mark.asyncio
    async def test_batch_query_with_jq_filter(self, batch_query_string, batch_variables_with_key):
        """Test batch query with jq filter applied."""
        jq_filter = ".data.target.approvedSymbol"

        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.side_effect = [
                QueryResult.create_success("TP53"),
                QueryResult.create_success("BRCA1"),
                QueryResult.create_success("BRCA2"),
            ]

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
                jq_filter=jq_filter,
            )

        # Verify jq_filter was passed to execute_graphql_query
        assert mock_execute.call_count == 3
        for call in mock_execute.call_args_list:
            # jq_filter is passed as a keyword argument
            assert call.kwargs["jq_filter"] == jq_filter

        assert isinstance(result, BatchQueryResult)
        assert result.summary.successful == 3

    @pytest.mark.asyncio
    async def test_batch_query_without_jq_filter(self, batch_query_string, batch_variables_with_key):
        """Test batch query without jq filter."""
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.side_effect = [
                QueryResult.create_success({"target": {"id": "ENSG00000141510"}}),
                QueryResult.create_success({"target": {"id": "ENSG00000012048"}}),
                QueryResult.create_success({"target": {"id": "ENSG00000139618"}}),
            ]

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
                jq_filter=None,
            )

        # Verify jq_filter was NOT passed (should be None)
        assert mock_execute.call_count == 3
        for call in mock_execute.call_args_list:
            # jq_filter is passed as a keyword argument
            assert call.kwargs.get("jq_filter") is None

        assert isinstance(result, BatchQueryResult)
        assert result.summary.successful == 3

    @pytest.mark.asyncio
    async def test_batch_query_exception_handling(self, batch_query_string, batch_variables_with_key):
        """Test that error results during individual query execution are handled."""
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.side_effect = [
                QueryResult.create_success({"target": {"id": "ENSG00000141510"}}),
                QueryResult.create_error("Network error"),
                QueryResult.create_success({"target": {"id": "ENSG00000139618"}}),
            ]

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
            )

        assert isinstance(result, BatchQueryResult)
        assert result.summary.total == 3
        assert result.summary.successful == 2
        assert result.summary.failed == 1

        # Check the error result
        result_dict = {r.key: r for r in result.results if r.key is not None}
        failed_result = result_dict["ENSG00000012048"]
        assert failed_result.result.status == QueryResultStatus.ERROR
        assert "Network error" in str(failed_result.result.message)

    @pytest.mark.asyncio
    async def test_batch_query_calls_execute_correctly(self, batch_query_string, batch_variables_with_key):
        """Test that batch query calls execute_graphql_query correctly."""
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.return_value = QueryResult.create_success({})

            await batch_query_fn(
                query_string=batch_query_string,
                variables_list=[batch_variables_with_key[0]],
                key_field="ensemblId",
            )

            # Verify the function was called with correct arguments
            mock_execute.assert_called_once_with(
                batch_query_string,
                batch_variables_with_key[0],
                jq_filter=None,
            )

    @pytest.mark.asyncio
    async def test_batch_query_parallel_execution(self, batch_query_string, batch_variables_with_key):
        """Test that queries are executed in parallel."""
        call_order = []

        async def track_call(query_string, variables, jq_filter=None):
            call_order.append(variables["ensemblId"])
            return QueryResult.create_success({})

        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            side_effect=track_call,
        ):
            await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
            )

        # Verify all queries were called (order may vary due to parallel execution)
        assert len(call_order) == 3
        assert set(call_order) == {"ENSG00000141510", "ENSG00000012048", "ENSG00000139618"}

    @pytest.mark.asyncio
    async def test_batch_query_result_structure(self, batch_query_string, batch_variables_with_key):
        """Test the structure of batch query results."""
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.return_value = QueryResult.create_success({"target": {"id": "test"}})

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=batch_variables_with_key,
                key_field="ensemblId",
            )

        # Check top-level structure
        assert isinstance(result, BatchQueryResult)
        assert isinstance(result.results, list)
        assert isinstance(result.summary, type(result.summary))

        # Check summary structure
        assert hasattr(result.summary, "total")
        assert hasattr(result.summary, "successful")
        assert hasattr(result.summary, "failed")
        assert hasattr(result.summary, "warning")

    @pytest.mark.asyncio
    async def test_batch_query_jq_filter_warning(self, batch_query_string, batch_variables_with_key):
        """Test that jq filter warnings are preserved in results."""
        with patch(
            "open_targets_platform_mcp.tools.batch_query.batch_query.execute_graphql_query",
            new_callable=AsyncMock,
        ) as mock_execute:
            mock_execute.return_value = QueryResult.create_warning(
                {"target": {"id": "test"}},
                "jq filter failed: null value",
            )

            result = await batch_query_fn(
                query_string=batch_query_string,
                variables_list=[batch_variables_with_key[0]],
                key_field="ensemblId",
                jq_filter=".data.target.missing",
            )

        # The warning should still be in the result
        result_dict = {r.key: r for r in result.results if r.key is not None}
        assert result_dict["ENSG00000141510"].result.status == QueryResultStatus.WARNING
        assert result.summary.warning == 1  # Counted as warning
        assert result.summary.successful == 0


# ============================================================================
# Integration Tests
# ============================================================================


@pytest.mark.integration
class TestBatchQueryIntegration:
    """Integration tests with real API calls."""

    @pytest.mark.asyncio
    async def test_real_batch_query(self):
        """Test real batch query against Open Targets Platform API."""
        query = """
        query GetTarget($ensemblId: String!) {
            target(ensemblId: $ensemblId) {
                id
                approvedSymbol
            }
        }
        """

        variables_list = [
            {"ensemblId": "ENSG00000141510"},
            {"ensemblId": "ENSG00000012048"},
        ]

        result = await batch_query_fn(query_string=query, variables_list=variables_list, key_field="ensemblId")

        assert isinstance(result, BatchQueryResult)
        assert result.summary.total == 2
        assert result.summary.successful == 2
        assert result.summary.failed == 0

        # Verify actual data
        result_dict = {r.key: r for r in result.results if r.key is not None}
        assert result_dict["ENSG00000141510"].result.result["target"]["approvedSymbol"] == "TP53"
        assert result_dict["ENSG00000012048"].result.result["target"]["approvedSymbol"] == "BRCA1"

    @pytest.mark.asyncio
    async def test_real_batch_query_with_jq_filter(self):
        """Test real batch query with jq filter (requires jq enabled)."""
        from open_targets_platform_mcp.settings import settings

        # Enable jq for this integration test
        original_jq_enabled = settings.jq_enabled
        settings.jq_enabled = True

        try:
            query = """
            query GetTarget($ensemblId: String!) {
                target(ensemblId: $ensemblId) {
                    id
                    approvedSymbol
                }
            }
            """

            variables_list = [
                {"ensemblId": "ENSG00000141510"},
                {"ensemblId": "ENSG00000012048"},
            ]

            result = await batch_query_fn(
                query_string=query,
                variables_list=variables_list,
                key_field="ensemblId",
                jq_filter=".target.approvedSymbol",
            )

            assert isinstance(result, BatchQueryResult)
            result_dict = {r.key: r for r in result.results if r.key is not None}
            # jq filter returns a list (even for single results)
            assert isinstance(result_dict["ENSG00000141510"].result.result, list)
            assert result_dict["ENSG00000141510"].result.result == ["TP53"]
            assert isinstance(result_dict["ENSG00000012048"].result.result, list)
            assert result_dict["ENSG00000012048"].result.result == ["BRCA1"]
        finally:
            # Restore original value
            settings.jq_enabled = original_jq_enabled
