"""
Tests for MCP Sub-Agent Tool Inheritance

These tests verify that sub-agents can properly inherit and use MCP tools
from their parent agents, addressing the 83% failure rate issue.
"""

import asyncio
import json
import os
from unittest.mock import MagicMock, Mock, patch

import pytest

from mcp_server.core.mcp_config_propagator import MCPConfigPropagator, MCPToolRegistryPropagator
from mcp_server.utils.sub_agent_helper import (
    SubAgentMCPHelper,
    SubAgentToolBridge,
    inherit_mcp_tools,
)


class TestMCPConfigPropagator:
    """Test MCP configuration propagation."""

    @pytest.fixture
    def sample_mcp_config(self, tmp_path):
        """Create a sample .mcp.json configuration."""
        config = {
            "mcpServers": {
                "code-index-mcp": {
                    "command": "node",
                    "args": ["/path/to/mcp-server.js"],
                    "env": {"INDEX_PATH": "/workspaces/.indexes", "LOG_LEVEL": "info"},
                }
            }
        }

        config_path = tmp_path / ".mcp.json"
        with open(config_path, "w") as f:
            json.dump(config, f)

        return config_path

    def test_config_loading(self, sample_mcp_config):
        """Test loading MCP configuration from file."""
        propagator = MCPConfigPropagator(sample_mcp_config)

        assert propagator.config is not None
        assert "mcpServers" in propagator.config
        assert "code-index-mcp" in propagator.config["mcpServers"]

    def test_environment_propagation(self, sample_mcp_config):
        """Test generation of propagation environment variables."""
        propagator = MCPConfigPropagator(sample_mcp_config)
        env_vars = propagator.get_propagation_env()

        # Check basic propagation flags
        assert env_vars["MCP_INHERIT_CONFIG"] == "true"
        assert env_vars["MCP_PROPAGATE_TOOLS"] == "true"
        assert env_vars["MCP_SUB_AGENT_ACCESS"] == "true"

        # Check server configuration propagation
        assert env_vars["MCP_SERVER_CODE-INDEX-MCP_COMMAND"] == "node"
        assert "MCP_SERVER_CODE-INDEX-MCP_ARGS" in env_vars
        assert env_vars["MCP_SERVER_CODE-INDEX-MCP_ENV_INDEX_PATH"] == "/workspaces/.indexes"

    def test_tool_registry_serialization(self):
        """Test serialization of tool registry."""
        propagator = MCPConfigPropagator()

        tools = [
            {
                "name": "symbol_lookup",
                "description": "Look up symbol definitions",
                "inputSchema": {
                    "type": "object",
                    "properties": {"symbol": {"type": "string"}},
                    "required": ["symbol"],
                },
            }
        ]

        serialized = propagator.serialize_tool_registry(tools)
        data = json.loads(serialized)

        assert data["version"] == "1.0"
        assert len(data["tools"]) == 1
        assert data["tools"][0]["name"] == "symbol_lookup"

    def test_environment_application(self, sample_mcp_config):
        """Test applying and restoring environment variables."""
        original_env = os.environ.copy()

        propagator = MCPConfigPropagator(sample_mcp_config)
        propagator.apply_to_environment()

        # Check environment was modified
        assert os.environ.get("MCP_INHERIT_CONFIG") == "true"
        assert "MCP_CONFIG_PATH" in os.environ

        # Restore environment
        propagator.restore_environment()

        # Check environment was restored
        for key, value in original_env.items():
            assert os.environ.get(key) == value

    def test_validation(self, sample_mcp_config):
        """Test configuration validation."""
        propagator = MCPConfigPropagator(sample_mcp_config)
        results = propagator.validate_propagation()

        assert results["config_found"] is True
        assert results["config_valid"] is True
        assert results["servers_configured"] is True
        assert results["env_vars_set"] is True
        assert results["inheritance_enabled"] is True


class TestSubAgentMCPHelper:
    """Test sub-agent MCP helper functionality."""

    @pytest.fixture
    def sub_agent_env(self):
        """Set up sub-agent environment."""
        original_env = os.environ.copy()

        # Set up sub-agent environment
        os.environ["MCP_SUB_AGENT_ACCESS"] = "true"
        os.environ["MCP_SERVER_CODE-INDEX-MCP_COMMAND"] = "node"
        os.environ["MCP_SERVER_CODE-INDEX-MCP_ARGS"] = '["server.js"]'
        os.environ["MCP_SERVER_CODE-INDEX-MCP_ENV_INDEX_PATH"] = "/indexes"

        # Set up tool registry
        tool_registry = {
            "version": "1.0",
            "tools": {
                "symbol_lookup": {
                    "description": "Look up symbols",
                    "inputSchema": {"required": ["symbol"]},
                },
                "search_code": {
                    "description": "Search code",
                    "inputSchema": {"required": ["query"]},
                },
            },
        }
        os.environ["MCP_TOOL_REGISTRY"] = json.dumps(tool_registry)

        yield

        # Restore original environment
        os.environ.clear()
        os.environ.update(original_env)

    def test_sub_agent_detection(self, sub_agent_env):
        """Test detection of sub-agent context."""
        helper = SubAgentMCPHelper()
        assert helper.is_sub_agent is True

    def test_inherited_config_loading(self, sub_agent_env):
        """Test loading inherited configuration from environment."""
        helper = SubAgentMCPHelper()

        assert "code-index-mcp" in helper.inherited_config["servers"]
        assert helper.inherited_config["servers"]["code-index-mcp"]["command"] == "node"

        assert "symbol_lookup" in helper.inherited_config["tools"]
        assert "search_code" in helper.inherited_config["tools"]

    @pytest.mark.asyncio
    async def test_tool_registration(self, sub_agent_env):
        """Test registration of inherited tools."""
        helper = SubAgentMCPHelper()
        tools = helper.register_inherited_tools()

        assert len(tools) == 2
        assert "symbol_lookup" in tools
        assert "search_code" in tools

        # Test tool execution
        result = await tools["symbol_lookup"](symbol="test_symbol")
        assert result["status"] == "success"
        assert result["tool"] == "symbol_lookup"

    @pytest.mark.asyncio
    async def test_tool_validation(self, sub_agent_env):
        """Test validation of required parameters."""
        helper = SubAgentMCPHelper()
        tools = helper.register_inherited_tools()

        # Should raise error for missing required parameter
        with pytest.raises(ValueError) as exc_info:
            await tools["symbol_lookup"]()

        assert "Missing required parameter: symbol" in str(exc_info.value)

    def test_tool_availability_check(self, sub_agent_env):
        """Test checking tool availability."""
        helper = SubAgentMCPHelper()
        helper.register_inherited_tools()

        availability = helper.validate_tool_availability()

        assert availability["symbol_lookup"] is True
        assert availability["search_code"] is True

        # Check for common MCP tools
        assert "mcp__code-index-mcp__symbol_lookup" in availability

    def test_mcp_server_command_retrieval(self, sub_agent_env):
        """Test retrieving MCP server command."""
        helper = SubAgentMCPHelper()

        command = helper.get_mcp_server_command("code-index-mcp")
        assert command == ["node", "server.js"]

        # Test non-existent server
        assert helper.get_mcp_server_command("non-existent") is None


class TestSubAgentToolBridge:
    """Test sub-agent tool bridge functionality."""

    @pytest.fixture
    def mock_helper(self):
        """Create a mock helper."""
        helper = Mock(spec=SubAgentMCPHelper)
        helper.get_mcp_server_command.return_value = ["node", "server.js"]
        helper.create_mcp_environment.return_value = os.environ.copy()
        return helper

    @pytest.mark.asyncio
    async def test_mcp_server_startup(self, mock_helper):
        """Test starting an MCP server."""
        bridge = SubAgentToolBridge(mock_helper)

        with patch("asyncio.create_subprocess_exec") as mock_create:
            mock_proc = MagicMock()
            mock_proc.returncode = None
            mock_create.return_value = mock_proc

            result = await bridge.ensure_mcp_server("test-server")
            assert result is True
            assert "test-server" in bridge.server_processes

    @pytest.mark.asyncio
    async def test_mcp_server_cleanup(self, mock_helper):
        """Test cleaning up MCP server processes."""
        bridge = SubAgentToolBridge(mock_helper)

        # Add a mock process
        mock_proc = MagicMock()
        mock_proc.returncode = None

        async def _wait():
            pass

        mock_proc.wait = _wait
        bridge.server_processes["test-server"] = mock_proc

        await bridge.cleanup()

        mock_proc.terminate.assert_called_once()


class TestMCPToolRegistryPropagator:
    """Test MCP tool registry propagation."""

    def test_tool_registration(self):
        """Test registering tools in the registry."""
        propagator = MCPToolRegistryPropagator()

        tool_def = {"description": "Test tool", "inputSchema": {"type": "object"}}

        propagator.register_tool("test_tool", tool_def)
        assert "test_tool" in propagator.tool_registry

    def test_registry_serialization(self):
        """Test serializing and deserializing the tool registry."""
        propagator = MCPToolRegistryPropagator()

        propagator.register_tool("tool1", {"description": "Tool 1"})
        propagator.register_tool("tool2", {"description": "Tool 2"})

        serialized = propagator.get_serialized_registry()
        data = json.loads(serialized)

        assert data["version"] == "1.0"
        assert data["count"] == 2
        assert "tool1" in data["tools"]
        assert "tool2" in data["tools"]

        # Test deserialization
        deserialized = propagator.deserialize_registry(serialized)
        assert len(deserialized) == 2
        assert "tool1" in deserialized


@pytest.mark.integration
class TestIntegration:
    """Integration tests for the full sub-agent flow."""

    @pytest.fixture
    def full_setup(self, tmp_path):
        """Set up a complete test environment."""
        # Create MCP config
        config = {
            "mcpServers": {
                "code-index-mcp": {
                    "command": "python",
                    "args": ["-m", "mcp_server"],
                    "env": {"INDEX_PATH": str(tmp_path / "indexes")},
                }
            }
        }

        config_path = tmp_path / ".mcp.json"
        with open(config_path, "w") as f:
            json.dump(config, f)

        return config_path

    def test_full_propagation_flow(self, full_setup):
        """Test the complete flow from parent to sub-agent."""
        # Parent agent setup
        propagator = MCPConfigPropagator(full_setup)
        tool_registry = MCPToolRegistryPropagator()

        # Register some tools
        tool_registry.register_tool(
            "symbol_lookup",
            {"description": "Look up symbol", "inputSchema": {"required": ["symbol"]}},
        )

        # Apply to environment
        propagator.apply_to_environment()
        tool_registry.apply_to_environment()

        # Sub-agent inherits tools
        helper = inherit_mcp_tools()

        assert helper is not None
        assert helper.is_sub_agent is True
        assert len(helper.tool_functions) > 0

        # Validate tools are available
        availability = helper.validate_tool_availability()
        assert availability["symbol_lookup"] is True

        # Clean up
        propagator.restore_environment()

    def test_task_agent_mcp_access(self, full_setup):
        """Test the specific scenario where Task agents couldn't access MCP tools."""
        # Simulate main agent setting up environment
        propagator = MCPConfigPropagator(full_setup)
        propagator.apply_to_environment()

        # Simulate Task agent trying to access MCP
        os.environ["CLAUDE_AGENT_TYPE"] = "task"  # Simulate task agent

        helper = inherit_mcp_tools()
        assert helper is not None

        # Check that MCP tools are available
        availability = helper.validate_tool_availability()
        mcp_tools = [k for k in availability.keys() if k.startswith("mcp__")]

        assert len(mcp_tools) > 0, "No MCP tools available in Task agent"

        # Clean up
        propagator.restore_environment()
        os.environ.pop("CLAUDE_AGENT_TYPE", None)


@pytest.mark.integration
class TestRealWorldScenarios:
    """Test real-world scenarios that were failing."""
