"""Tests for the Graph MCP client implementation."""

from datetime import UTC
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from robosystems.middleware.mcp import (
  GraphAPIError,
  GraphMCPClient,
  GraphMCPTools,
)


@pytest.fixture
def mock_async_graph_client():
  """Mock the GraphClient."""
  with patch("robosystems.middleware.mcp.client.GraphClient") as mock:
    client_instance = AsyncMock()
    mock.return_value = client_instance
    yield client_instance


@pytest.fixture
def mock_httpx_client():
  """Mock the httpx.AsyncClient."""
  with patch("robosystems.middleware.mcp.client.httpx.AsyncClient") as mock:
    client_instance = AsyncMock()
    mock.return_value = client_instance
    yield client_instance


class TestGraphMCPClient:
  """Unit tests for the GraphMCPClient class."""

  @pytest.mark.unit
  def test_init_with_api_url(self):
    """Test initialization with API URL."""
    with (
      patch("robosystems.middleware.mcp.client.GraphClient"),
      patch("robosystems.middleware.mcp.client.httpx.AsyncClient"),
    ):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      assert client.api_base_url == "http://test:8001"
      assert client.graph_id == "test"
      assert client.timeout == 30  # Uses TuningConfig.get_graph_http_timeout() default

  @pytest.mark.unit
  def test_init_custom_timeout(self):
    """Test initialization with custom timeout."""
    with (
      patch("robosystems.middleware.mcp.client.GraphClient"),
      patch("robosystems.middleware.mcp.client.httpx.AsyncClient"),
    ):
      client = GraphMCPClient(
        api_base_url="http://test:8001", graph_id="test", timeout=60
      )
      assert client.timeout == 60

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_execute_query_success(self, mock_async_graph_client):
    """Test successful query execution."""
    # Mock successful response from GraphClient
    mock_async_graph_client.query.return_value = {
      "data": [{"name": "Test Entity", "cik": "0000123456"}],
      "execution_time_ms": 15.5,
    }

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      result = await client.execute_query(
        "MATCH (c:Entity) RETURN c.name as name, c.cik as cik LIMIT 1"
      )

      assert len(result) == 1
      assert result[0]["name"] == "Test Entity"
      assert result[0]["cik"] == "0000123456"

      # Verify the correct API call was made
      # Note: Auto-LIMIT feature will modify queries without explicit LIMIT
      mock_async_graph_client.query.assert_called_once_with(
        cypher="MATCH (c:Entity) RETURN c.name as name, c.cik as cik LIMIT 1",
        graph_id="test",
        parameters=None,
      )

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_execute_query_with_parameters(self, mock_async_graph_client):
    """Test query execution with parameters."""
    mock_async_graph_client.query.return_value = {
      "data": [{"name": "Apple Inc", "cik": "0000320193"}],
      "execution_time_ms": 12.3,
    }

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      result = await client.execute_query(
        "MATCH (c:Entity {cik: $cik}) RETURN c.name as name, c.cik as cik",
        {"cik": "0000320193"},
      )

      assert len(result) == 1
      assert result[0]["name"] == "Apple Inc"

      # Verify parameters were passed
      # Note: Auto-LIMIT of 1000 is appended to queries without explicit LIMIT
      mock_async_graph_client.query.assert_called_once_with(
        cypher="MATCH (c:Entity {cik: $cik}) RETURN c.name as name, c.cik as cik LIMIT 1000",
        graph_id="test",
        parameters={"cik": "0000320193"},
      )

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_execute_query_http_error(self, mock_async_graph_client):
    """Test query execution with HTTP error."""
    from httpx import HTTPError

    mock_response = MagicMock()
    mock_response.status_code = 400
    mock_response.text = "Bad Request"
    mock_response.json.side_effect = Exception(
      "JSON decode error"
    )  # Simulate JSON parse failure

    http_error = HTTPError("HTTP Error")
    http_error.response = mock_response
    mock_async_graph_client.query.side_effect = http_error

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      with pytest.raises(GraphAPIError, match="Invalid query"):
        await client.execute_query("INVALID QUERY")

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_get_schema(self, mock_async_graph_client):
    """Test schema retrieval (no counts for performance)."""
    # Mock responses - schema no longer fetches counts
    query_responses = [
      # SHOW_TABLES response only
      [
        {"name": "Entity", "type": "NODE"},
        {"name": "Report", "type": "NODE"},
        {"name": "HAS_REPORT", "type": "REL"},
      ],
    ]

    # Set up the GraphClient mock to return different responses for each call
    mock_async_graph_client.query.side_effect = [
      {"data": response, "execution_time_ms": 10} for response in query_responses
    ]

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      schema = await client.get_schema()

      assert len(schema) == 3
      # Check node tables
      entity_schema = next(s for s in schema if s["label"] == "Entity")
      assert entity_schema["type"] == "node"
      assert "description" in entity_schema
      assert "count" not in entity_schema  # No counts for performance

      # Check relationship table
      rel_schema = next(s for s in schema if s["label"] == "HAS_REPORT")
      assert rel_schema["type"] == "relationship"
      assert "from_node" in rel_schema
      assert "to_node" in rel_schema
      assert "count" not in rel_schema  # No counts for performance

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_get_graph_info(self, mock_async_graph_client, mock_httpx_client):
    """Test graph info retrieval (optimized, no relationship counts)."""
    # Mock API info response
    info_response = MagicMock()
    info_response.raise_for_status = MagicMock()
    info_response.json.return_value = {
      "database_path": "/path/to/db",
      "read_only": True,
      "uptime_seconds": 3600,
    }

    # Mock the queries responses - no relationship counts for performance
    query_responses = [
      # SHOW_TABLES response
      [
        {"name": "Entity", "type": "NODE"},
        {"name": "Report", "type": "NODE"},
        {"name": "HAS_REPORT", "type": "REL"},
      ],
      # Aggregated node counts: MATCH (n) RETURN labels(n) as label, count(*) as cnt
      [
        {"label": "Entity", "cnt": 100},
        {"label": "Report", "cnt": 50},
      ],
    ]

    mock_httpx_client.get.return_value = info_response
    mock_async_graph_client.query.side_effect = [
      {"data": response, "execution_time_ms": 10} for response in query_responses
    ]

    with patch(
      "robosystems.middleware.mcp.client.httpx.AsyncClient",
      return_value=mock_httpx_client,
    ):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      info = await client.get_graph_info()

      assert info["graph_id"] == "test"
      assert info["total_nodes"] == 150  # Sum of Entity (100) + Report (50)
      assert "total_relationships" not in info  # No longer included
      assert "relationship_types" in info  # Just the labels, not counts
      assert info["read_only"] is True
      assert "Entity" in info["node_labels"]
      assert "Report" in info["node_labels"]
      assert "HAS_REPORT" in info["relationship_types"]

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_close(self, mock_async_graph_client, mock_httpx_client):
    """Test client close."""
    with patch(
      "robosystems.middleware.mcp.client.httpx.AsyncClient",
      return_value=mock_httpx_client,
    ):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      await client.close()

      mock_async_graph_client.close.assert_called_once()
      mock_httpx_client.aclose.assert_called_once()


class TestGraphMCPTools:
  """Unit tests for the GraphMCPTools class."""

  @pytest.fixture
  def mock_graph_client(self):
    """Create mock GraphMCPClient."""
    return AsyncMock(spec=GraphMCPClient)

  @pytest.mark.unit
  def test_get_tool_definitions_core_only(self, mock_graph_client):
    """Test tool definitions for a graph with no schema extensions (custom graph)."""
    mock_graph_client.graph_id = "kg_custom_graph"
    tools = GraphMCPTools(mock_graph_client, schema_extensions=[])
    definitions = tools.get_tool_definitions_as_dict()

    tool_names = {t["name"] for t in definitions}

    # Core tools (2) should always be present
    assert "read-graph-cypher" in tool_names
    assert "get-graph-schema" in tool_names

    # roboledger extension tools should NOT be present
    assert "get-example-queries" not in tool_names
    assert "financial-statement-analysis" not in tool_names
    assert "live-financial-statement" not in tool_names
    assert "resolve-element" not in tool_names
    assert "resolve-structure" not in tool_names

    # Removed tools should NOT be present
    assert "discover-properties" not in tool_names
    assert "list-disclosures" not in tool_names
    assert "get-disclosure-detail" not in tool_names

    # Verify core tool content
    cypher_tool = next(t for t in definitions if t["name"] == "read-graph-cypher")
    assert "Execute read-only Cypher queries" in cypher_tool["description"]
    assert "query" in cypher_tool["inputSchema"]["properties"]

    schema_tool = next(t for t in definitions if t["name"] == "get-graph-schema")
    assert "Get the complete database schema" in schema_tool["description"]

  @pytest.mark.unit
  def test_get_tool_definitions_with_roboledger(self, mock_graph_client):
    """Test tool definitions for a graph with roboledger schema extension."""
    mock_graph_client.graph_id = "test_graph"
    tools = GraphMCPTools(mock_graph_client, schema_extensions=["roboledger"])
    definitions = tools.get_tool_definitions_as_dict()

    tool_names = {t["name"] for t in definitions}

    # Core tools should be present
    assert "read-graph-cypher" in tool_names
    assert "get-graph-schema" in tool_names

    # roboledger extension tools should be present
    assert "get-example-queries" in tool_names
    assert "financial-statement-analysis" in tool_names
    # live-financial-statement is only registered on non-shared, read-write
    # graphs — by default mock_graph_client isn't a shared repo so it's
    # expected here.
    assert "live-financial-statement" in tool_names

    # Removed tools should NOT be present
    assert "discover-properties" not in tool_names
    assert "list-disclosures" not in tool_names
    assert "get-disclosure-detail" not in tool_names

    # Verify example queries tool content
    example_tool = next(t for t in definitions if t["name"] == "get-example-queries")
    assert "Get example Cypher queries" in example_tool["description"]
    assert "category" in example_tool["inputSchema"]["properties"]

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_call_tool_cypher_success(self, mock_graph_client):
    """Test successful Cypher tool execution."""
    mock_graph_client.graph_id = "test_graph"  # Add missing graph_id attribute
    mock_graph_client.execute_query.return_value = [
      {"name": "Apple Inc", "cik": "0000320193"}
    ]

    tools = GraphMCPTools(mock_graph_client)
    result = await tools.call_tool(
      "read-graph-cypher",
      {"query": "MATCH (c:Entity) RETURN c.name as name, c.cik as cik LIMIT 1"},
      return_raw=True,
    )

    assert len(result) == 1
    assert result[0]["name"] == "Apple Inc"
    mock_graph_client.execute_query.assert_called_once()

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_call_tool_cypher_write_blocked(self, mock_graph_client):
    """Test that write operations are blocked."""
    mock_graph_client.graph_id = "test_graph"  # Add missing graph_id attribute
    # Add the _is_read_only_query method to the mock
    mock_graph_client._is_read_only_query = MagicMock(return_value=False)

    tools = GraphMCPTools(mock_graph_client)

    result = await tools.call_tool(
      "read-graph-cypher", {"query": "CREATE (c:Entity {name: 'Test'})"}
    )

    # The tool catches the error and returns an error message
    assert "Only read-only queries are allowed" in result

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_call_tool_schema(self, mock_graph_client):
    """Test schema tool execution."""
    mock_schema = [{"label": "Entity", "type": "node", "properties": []}]
    mock_graph_client.get_schema.return_value = mock_schema

    tools = GraphMCPTools(mock_graph_client)
    result = await tools.call_tool("get-graph-schema", {}, return_raw=True)

    assert result == mock_schema
    mock_graph_client.get_schema.assert_called_once()

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_call_tool_unknown(self, mock_graph_client):
    """Test unknown tool handling."""
    tools = GraphMCPTools(mock_graph_client)

    result = await tools.call_tool("unknown-tool", {})
    assert "Unknown tool" in result

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_execute_cypher_tool(self, mock_graph_client):
    """Test Cypher tool wrapper method."""
    mock_graph_client.graph_id = "test_graph"  # Add missing graph_id attribute
    mock_graph_client.execute_query.return_value = [{"count": 5}]

    tools = GraphMCPTools(mock_graph_client)
    result = await tools.execute_cypher_tool(
      "MATCH (c:Entity) RETURN count(c) as count"
    )

    assert result[0]["count"] == 5

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_execute_schema_tool(self, mock_graph_client):
    """Test schema tool wrapper method."""
    mock_schema = [{"label": "Entity", "type": "node"}]
    mock_graph_client.get_schema.return_value = mock_schema

    tools = GraphMCPTools(mock_graph_client)
    result = await tools.execute_schema_tool()

    assert result == mock_schema

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_call_workspace_tools(self, mock_graph_client, monkeypatch):
    """Test routing calls to workspace tools."""
    from datetime import datetime

    # Enable workspace tools for this test
    monkeypatch.setattr("robosystems.config.env.MCP_WORKSPACE_ENABLED", True)

    mock_graph_client.graph_id = "kg1234567890abcdef"
    mock_graph_client.user = MagicMock()
    mock_graph_client.user.id = "user123"

    tools = GraphMCPTools(mock_graph_client)

    with patch("robosystems.database.get_db_session") as mock_db:
      mock_session = MagicMock()
      mock_db.return_value = iter([mock_session])

      mock_graph = MagicMock()
      mock_graph.graph_id = "kg1234567890abcdef"
      mock_graph.graph_name = "Test Graph"
      mock_graph.parent_graph_id = None
      mock_graph.created_at = datetime(2025, 1, 1, tzinfo=UTC)

      # Mock the query chain for both subgraphs and parent graph
      mock_query = MagicMock()
      mock_filter = MagicMock()
      mock_order_by = MagicMock()

      mock_session.query.return_value = mock_query
      mock_query.filter.return_value = mock_filter
      mock_filter.order_by.return_value = mock_order_by
      mock_filter.first.return_value = mock_graph
      mock_order_by.all.return_value = []

      result = await tools.call_tool("list-subgraphs", {}, return_raw=True)

    # Check that the tool was called and returned a result
    assert result is not None
    assert isinstance(result, dict)
    # list-subgraphs returns subgraphs array or error
    assert "subgraphs" in result or "error" in result
    if "subgraphs" in result:
      # Primary graph + any subgraphs
      assert result["total_subgraphs"] == 0
      assert len(result["subgraphs"]) == 1  # primary entry only

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_call_data_operation_tools(self, mock_graph_client, monkeypatch):
    """Test routing calls to data operation tools.

    build-fact-grid now delegates to the ops-layer ``query_fact_grid``
    (see ``operations/roboledger/views/fact_query.py``) — patching there
    keeps the test focused on the manager's dispatch contract.
    """
    monkeypatch.setattr("robosystems.config.env.FACT_GRID_ENABLED", True)

    mock_graph_client.graph_id = "kg1234567890abcdef"

    tools = GraphMCPTools(mock_graph_client)

    import pandas as pd

    fact_df = pd.DataFrame(
      [
        {
          "element_id": "test:Element",
          "element_name": "Cash",
          "period_end": "2025-01-01",
          "value": 1000.0,
          "unit": "USD",
        }
      ]
    )

    with patch(
      "robosystems.middleware.mcp.tools.fact_grid_tool.query_fact_grid",
      new_callable=AsyncMock,
      return_value=fact_df,
    ):
      result = await tools.call_tool(
        "build-fact-grid",
        {
          "elements": ["test:Element"],
          "periods": ["2025-01-01"],
        },
        return_raw=True,
      )

    assert result["success"] is True
    assert result["fact_count"] == 1

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_workspace_tool_error_handling(self, mock_graph_client, monkeypatch):
    """Test error handling for workspace tools."""
    # Enable workspace tools for this test
    monkeypatch.setattr("robosystems.config.env.MCP_WORKSPACE_ENABLED", True)

    mock_graph_client.graph_id = "kg1234567890abcdef"
    mock_graph_client.user = MagicMock()

    tools = GraphMCPTools(mock_graph_client)

    result = await tools.call_tool(
      "create-subgraph", {"name": "invalid-name"}, return_raw=True
    )

    assert result["error"] == "invalid_name"

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_data_tool_error_handling(self, mock_graph_client, monkeypatch):
    """Test error handling for data operation tools."""
    # Enable fact grid tool for this test
    monkeypatch.setattr("robosystems.config.env.FACT_GRID_ENABLED", True)

    mock_graph_client.graph_id = "kg1234567890abcdef"

    tools = GraphMCPTools(mock_graph_client)

    result = await tools.call_tool("build-fact-grid", {}, return_raw=True)

    assert result["error"] == "missing_elements"

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_workspace_tool_disabled_returns_error(
    self, mock_graph_client, monkeypatch
  ):
    """Test that workspace tools return proper error when disabled."""
    # Disable workspace tools for this test
    monkeypatch.setattr("robosystems.config.env.MCP_WORKSPACE_ENABLED", False)

    mock_graph_client.graph_id = "kg1234567890abcdef"
    tools = GraphMCPTools(mock_graph_client)

    # Calling a subgraph write tool when disabled returns an error envelope.
    result = await tools.call_tool(
      "create-subgraph", {"name": "test"}, return_raw=False
    )

    assert "not available" in result
    assert "MCP_WORKSPACE_ENABLED" in result

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_fact_grid_tool_disabled_returns_error(
    self, mock_graph_client, monkeypatch
  ):
    """Test that fact grid tool returns proper error when disabled."""
    # Disable fact grid tool for this test
    monkeypatch.setattr("robosystems.config.env.FACT_GRID_ENABLED", False)

    mock_graph_client.graph_id = "kg1234567890abcdef"
    tools = GraphMCPTools(mock_graph_client)

    result = await tools.call_tool(
      "build-fact-grid",
      {"elements": ["test"], "periods": ["2025-01-01"]},
      return_raw=False,
    )

    assert "not available" in result
    assert "FACT_GRID_ENABLED" in result

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_roboledger_tool_disabled_returns_error(self, mock_graph_client):
    """Test that roboledger-gated tools return proper error on custom graphs."""
    mock_graph_client.graph_id = "kg_custom_graph"
    tools = GraphMCPTools(mock_graph_client, schema_extensions=[])

    for tool_name in [
      "financial-statement-analysis",
      "live-financial-statement",
      "get-example-queries",
    ]:
      result = await tools.call_tool(tool_name, {}, return_raw=False)
      assert "not available" in result, f"{tool_name} should not be available"
      assert "roboledger" in result, f"{tool_name} error should mention roboledger"


class TestGraphMCPConfigurableSchema:
  """Test configurable schema tables."""

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_schema_all_tables_listed(self, mock_async_graph_client, monkeypatch):
    """Test that all tables are listed in schema (no counts for performance)."""

    # Mock table response - includes nodes and relationships
    tables_response = [
      {"name": "CustomTable1", "type": "NODE"},
      {"name": "CustomTable2", "type": "NODE"},
      {"name": "OtherTable", "type": "NODE"},
      {"name": "HAS_CUSTOM", "type": "REL"},
    ]

    mock_async_graph_client.query.side_effect = [
      {"data": tables_response, "execution_time_ms": 10},  # SHOW_TABLES only
    ]

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      schema = await client.get_schema()

      # Verify all tables are listed (3 nodes + 1 relationship)
      assert len(schema) == 4

      # Check that all node tables are present (no counts)
      custom1 = next(s for s in schema if s["label"] == "CustomTable1")
      assert custom1["type"] == "node"
      assert "count" not in custom1

      custom2 = next(s for s in schema if s["label"] == "CustomTable2")
      assert custom2["type"] == "node"

      other = next(s for s in schema if s["label"] == "OtherTable")
      assert other["type"] == "node"

      # Check relationship is listed
      has_custom = next(s for s in schema if s["label"] == "HAS_CUSTOM")
      assert has_custom["type"] == "relationship"
      assert "from_node" in has_custom
      assert "to_node" in has_custom

      # Verify only 1 query (SHOW_TABLES)
      query_calls = list(mock_async_graph_client.query.call_args_list)
      assert len(query_calls) == 1


class TestGraphMCPFactory:
  """Test the factory function for creating Graph MCP clients."""

  @pytest.mark.unit
  @pytest.mark.asyncio
  async def test_create_graph_mcp_client_default(self):
    """Test factory with default parameters."""
    with (
      patch("robosystems.middleware.mcp.client.GraphClient"),
      patch("robosystems.middleware.mcp.client.httpx.AsyncClient"),
      patch(
        "robosystems.graph_api.client.factory.GraphClientFactory.create_client"
      ) as mock_factory,
      patch(
        "robosystems.middleware.graph.utils.MultiTenantUtils.is_shared_repository",
        return_value=True,
      ),
    ):
      # Mock the factory to return a client with config.base_url attribute
      mock_graph_client = MagicMock()
      mock_graph_client.config.base_url = "http://localhost:8001"
      mock_graph_client.close = AsyncMock()

      # Mock factory returns the mock client directly (not a coroutine)
      async def create_client_async(*args, **kwargs):
        return mock_graph_client

      mock_factory.side_effect = create_client_async

      with patch.dict("os.environ", {"ENVIRONMENT": "dev"}):
        from robosystems.middleware.mcp import create_graph_mcp_client

        client = await create_graph_mcp_client()
        assert client.graph_id == "sec"
        assert "localhost" in client.api_base_url

  @pytest.mark.unit
  @pytest.mark.asyncio
  async def test_create_graph_mcp_client_prod(self):
    """Test factory with production environment."""
    with (
      patch("robosystems.middleware.mcp.client.GraphClient"),
      patch("robosystems.middleware.mcp.client.httpx.AsyncClient"),
      patch(
        "robosystems.graph_api.client.factory.GraphClientFactory.create_client"
      ) as mock_factory,
      patch(
        "robosystems.middleware.graph.utils.MultiTenantUtils.is_shared_repository",
        return_value=True,
      ),
    ):
      # Mock the factory to return a client with production config.base_url
      mock_graph_client = MagicMock()
      mock_graph_client.config.base_url = "http://robosystems.internal:8001"
      mock_graph_client.close = AsyncMock()

      # Mock factory returns the mock client directly (not a coroutine)
      async def create_client_async(*args, **kwargs):
        return mock_graph_client

      mock_factory.side_effect = create_client_async

      with patch(
        "robosystems.middleware.mcp.factory.env.is_production", return_value=True
      ):
        from robosystems.middleware.mcp import create_graph_mcp_client

        client = await create_graph_mcp_client()
        assert "robosystems.internal" in client.api_base_url

  @pytest.mark.unit
  @pytest.mark.asyncio
  async def test_create_graph_mcp_client_custom_url(self):
    """Test factory with custom URL override."""
    with (
      patch("robosystems.middleware.mcp.client.GraphClient"),
      patch("robosystems.middleware.mcp.client.httpx.AsyncClient"),
    ):
      from robosystems.middleware.mcp import create_graph_mcp_client

      client = await create_graph_mcp_client(
        api_base_url="http://custom:8001", graph_id="custom"
      )
      assert client.api_base_url == "http://custom:8001"
      assert client.graph_id == "custom"


class TestGraphMCPErrorSanitization:
  """Test error message sanitization for security."""

  @pytest.mark.unit
  def test_sanitize_error_removes_sensitive_info(self):
    """Test that sensitive information is removed from errors."""
    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")

      # Test various error types with sensitive information
      test_cases = [
        # File paths
        (
          Exception("Error at /home/user/robosystems/mcp.py line 123"),
          "Error at [REDACTED] [REDACTED]",
        ),
        # Memory addresses
        (
          Exception("Segfault at 0x7fff8b4c2340 in function"),
          "Segfault at [REDACTED] in function",
        ),
        # IP addresses
        (
          Exception("Cannot connect to 192.168.1.100 port 8001"),
          "Cannot connect to [REDACTED] [REDACTED]",
        ),
        # Database files
        (
          Exception("Cannot open database file test_entity.lbug"),
          "Cannot open database file [REDACTED]",
        ),
      ]

      for error, expected in test_cases:
        result = client._sanitize_error_message(error)
        assert result == expected, f"Expected '{expected}', got '{result}'"

  @pytest.mark.unit
  def test_sanitize_error_maps_common_errors(self):
    """Test that common errors are mapped to user-friendly messages."""
    from robosystems.config import env

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")

      # Determine expected query error message based on environment
      # In production, query errors get generic messages for security
      # In dev/staging, they're preserved for debugging
      if env.ENVIRONMENT in ("dev", "staging"):
        query_error_expected = "Syntax error near 'MATCH' at position 15"
      else:
        query_error_expected = (
          "Query validation failed. Please check your query syntax."
        )

      test_cases = [
        # Connection errors
        (
          ConnectionError("Connection refused to localhost:8001"),
          "Service temporarily unavailable. Please try again later.",
        ),
        # Timeout
        (
          TimeoutError("Request timed out after 30s"),
          "Request timed out during operation. Try a simpler query or increase timeout.",
        ),
        # Auth errors
        (
          Exception("Unauthorized: Invalid API key abc123"),
          "Authentication required. Please check your credentials.",
        ),
        # Query errors (environment-dependent handling)
        (
          Exception("Syntax error near 'MATCH' at position 15"),
          query_error_expected,
        ),
        # Resource errors
        (
          Exception("Out of memory while processing query"),
          "Query requires too many resources. Try limiting results or simplifying the query.",
        ),
      ]

      for error, expected in test_cases:
        result = client._sanitize_error_message(error)
        assert result == expected, (
          f"For error '{error}', expected '{expected}', got '{result}'"
        )

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_http_error_sanitization(
    self, mock_async_graph_client, mock_httpx_client
  ):
    """Test that HTTP errors are properly sanitized."""
    from httpx import HTTPError

    # Create mock response with sensitive details
    mock_response = MagicMock()
    mock_response.status_code = 500
    mock_response.json.return_value = {
      "detail": "Database error at /var/lib/lbug/test.db: disk full"
    }

    http_error = HTTPError("Internal server error")
    http_error.response = mock_response
    mock_async_graph_client.query.side_effect = http_error

    with patch(
      "robosystems.middleware.mcp.client.httpx.AsyncClient",
      return_value=mock_httpx_client,
    ):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      with pytest.raises(GraphAPIError) as exc_info:
        await client.execute_query("MATCH (n) RETURN n")

      # Error should be sanitized - sensitive path should NOT be in the error
      error_msg = str(exc_info.value)
      assert "/var/lib/lbug/" not in error_msg
      assert "test.db" not in error_msg
      # Should get generic 500 error message instead
      assert "server error" in error_msg.lower()


class TestGraphMCPReadOnlyValidation:
  """Test the enhanced read-only query validation."""

  @pytest.fixture
  def mock_async_graph_client(self):
    """Create a mock GraphClient."""
    mock = AsyncMock()
    return mock

  @pytest.mark.unit
  def test_read_only_validation_allows_valid_queries(self):
    """Test that valid read-only queries are allowed."""
    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")

      valid_queries = [
        "MATCH (n) RETURN n",
        "MATCH (c:Entity) WHERE c.name = 'Apple' RETURN c",
        "MATCH (c:Entity)-[:HAS_REPORT]->(r:Report) RETURN c, r",
        "MATCH (n) WHERE n.created_at > date('2024-01-01') RETURN n",
        "WITH 'test' as value MATCH (n) WHERE n.name = value RETURN n",
        "MATCH (n) RETURN n ORDER BY n.created_at DESC LIMIT 10",
        "MATCH (n) RETURN count(n) as total",
        "CALL SHOW_TABLES() RETURN name, type",
      ]

      for query in valid_queries:
        # Should not raise an exception
        assert client._is_read_only_query(query), f"Query should be valid: {query}"

  @pytest.mark.unit
  def test_read_only_validation_blocks_write_queries(self):
    """Test that write operations are blocked."""
    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")

      write_queries = [
        "CREATE (n:Person {name: 'John'})",
        "MATCH (n) SET n.name = 'NewName'",
        "MATCH (n) DELETE n",
        "MATCH (n) REMOVE n.property",
        "MERGE (n:Entity {name: 'Test'})",
        "DROP INDEX idx_name",
        "CREATE INDEX idx_name FOR (n:Person) ON (n.name)",
        "MATCH (a), (b) CREATE (a)-[:KNOWS]->(b)",
        "CALL db.createLabel('NewLabel')",
        "CALL apoc.refactor.rename.label('Old', 'New')",
      ]

      for query in write_queries:
        assert not client._is_read_only_query(query), (
          f"Query should be blocked: {query}"
        )

  @pytest.mark.unit
  def test_read_only_validation_handles_edge_cases(self):
    """Test edge cases and tricky patterns."""
    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")

      # These should be allowed (field names that look like keywords)
      allowed_edge_cases = [
        "MATCH (n) WHERE n.created_at > date('2024-01-01') RETURN n",
        "MATCH (n) WHERE n.set_name = 'test' RETURN n",
        "MATCH (n) WHERE n.delete_flag = false RETURN n",
        "MATCH (n) RETURN n.created_at, n.updated_at",
      ]

      for query in allowed_edge_cases:
        assert client._is_read_only_query(query), f"Query should be valid: {query}"

      # These should be blocked (sneaky write attempts)
      blocked_edge_cases = [
        "MATCH (n) SET n.value = 123",  # SET with property
        "match (n) create (m:Node) return n, m",  # lowercase keywords
        "MATCH (n) /* comment */ DELETE n",  # with comments
        "MATCH (a), (b) WHERE a.id = 1 AND b.id = 2 CREATE (a)-[:KNOWS]->(b)",
      ]

      for query in blocked_edge_cases:
        assert not client._is_read_only_query(query), (
          f"Query should be blocked: {query}"
        )


class TestGraphMCPAutoLimit:
  """Test the auto-LIMIT functionality for MCP context safety."""

  @pytest.fixture
  def mock_async_graph_client(self):
    """Create a mock GraphClient."""
    mock = AsyncMock()
    return mock

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_auto_limit_appended(self, mock_async_graph_client):
    """Test that LIMIT is automatically appended to queries without one."""
    mock_async_graph_client.query.return_value = {
      "data": [{"id": i} for i in range(1000)],
      "execution_time_ms": 50,
    }

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      # Query without LIMIT should get one added
      result = await client.execute_query("MATCH (n) RETURN n")

      # Verify LIMIT was appended
      mock_async_graph_client.query.assert_called_once()
      actual_query = mock_async_graph_client.query.call_args[1]["cypher"]
      assert "LIMIT 1000" in actual_query

      # Check for truncation marker
      assert len(result) == 1001  # 1000 results + 1 truncation marker
      assert result[-1]["_mcp_note"] == "RESULTS_TRUNCATED"

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_existing_limit_preserved(self, mock_async_graph_client):
    """Test that queries with existing LIMIT are not modified."""
    mock_async_graph_client.query.return_value = {
      "data": [{"id": i} for i in range(50)],
      "execution_time_ms": 20,
    }

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      # Query with existing LIMIT should not be modified
      original_query = "MATCH (n) RETURN n LIMIT 50"
      result = await client.execute_query(original_query)

      # Verify query was not modified
      mock_async_graph_client.query.assert_called_once_with(
        cypher=original_query,
        graph_id="test",
        parameters=None,
      )

      # No truncation marker for queries with explicit LIMIT
      assert len(result) == 50
      assert not any("_mcp_note" in row for row in result)

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_auto_limit_disabled(self, mock_async_graph_client, monkeypatch):
    """Test that auto-limit can be disabled via environment variable."""
    monkeypatch.setattr("robosystems.config.env.MCP_AUTO_LIMIT_ENABLED", False)

    mock_async_graph_client.query.return_value = {
      "data": [{"id": i} for i in range(100)],
      "execution_time_ms": 30,
    }

    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      original_query = "MATCH (n) RETURN n"
      result = await client.execute_query(original_query)

      # Verify query was not modified when auto-limit is disabled
      mock_async_graph_client.query.assert_called_once_with(
        cypher=original_query,
        graph_id="test",
        parameters=None,
      )

      assert len(result) == 100

  @pytest.mark.asyncio
  @pytest.mark.unit
  async def test_size_based_truncation(self, mock_async_graph_client):
    """Test that results are truncated when exceeding size limit."""
    # Create large result set
    large_data = [
      {"id": i, "data": "x" * 10000}  # Each row is ~10KB
      for i in range(100)
    ]

    mock_async_graph_client.query.return_value = {
      "data": large_data,
      "execution_time_ms": 100,
    }

    with (
      patch("robosystems.middleware.mcp.client.httpx.AsyncClient"),
      patch(
        "robosystems.middleware.mcp.client.TuningConfig.get_mcp_max_result_size_mb",
        return_value=0.001,  # 1KB limit
      ),
    ):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")
      client.graph_client = mock_async_graph_client

      result = await client.execute_query("MATCH (n) RETURN n")

      # Results should be truncated due to size limit
      assert len(result) < len(large_data)
      assert result[-1]["_mcp_note"] == "RESULTS_TRUNCATED_BY_SIZE"
      assert "_mcp_size_mb" in result[-1]

  @pytest.mark.unit
  def test_intelligent_limit_injection(self):
    """Test intelligent LIMIT injection for complex queries."""
    with patch("robosystems.middleware.mcp.client.httpx.AsyncClient"):
      client = GraphMCPClient(api_base_url="http://test:8001", graph_id="test")

      # Test simple query
      result = client._inject_limit_intelligently("MATCH (n) RETURN n", 100)
      assert result == "MATCH (n) RETURN n LIMIT 100"

      # Test query with ORDER BY
      result = client._inject_limit_intelligently(
        "MATCH (n) RETURN n ORDER BY n.name DESC", 100
      )
      assert result == "MATCH (n) RETURN n ORDER BY n.name DESC LIMIT 100"

      # Test UNION query
      result = client._inject_limit_intelligently(
        "MATCH (a:TypeA) RETURN a.name UNION MATCH (b:TypeB) RETURN b.name", 100
      )
      assert "LIMIT 100" in result
      # Both parts should have LIMIT
      parts = result.split("UNION")
      assert all("LIMIT 100" in part for part in parts)

      # Test query with semicolon
      result = client._inject_limit_intelligently("MATCH (n) RETURN n;", 100)
      assert result == "MATCH (n) RETURN n LIMIT 100"

      # Test query that already has LIMIT
      original = "MATCH (n) RETURN n LIMIT 50"
      result = client._inject_limit_intelligently(original, 100)
      assert result == original  # Should not change

      # Test complex query with WITH clause
      result = client._inject_limit_intelligently(
        "MATCH (n) WITH n, count(*) as cnt WHERE cnt > 5 RETURN n ORDER BY cnt", 100
      )
      assert (
        result
        == "MATCH (n) WITH n, count(*) as cnt WHERE cnt > 5 RETURN n ORDER BY cnt LIMIT 100"
      )
