"""
Contract tests for src_mcp_proxy.rs - MCP Proxy Server

This test suite validates the MCP Proxy Server implementation against its contract.
Tests cover configuration parsing, tool routing, policy enforcement, upstream connections,
and all documented error cases and invariants.

Test organization:
- Configuration and deserialization tests
- Tool listing and aggregation tests
- Tool calling and routing tests
- Policy enforcement tests
- Upstream connection lifecycle tests
- Proxy server lifecycle tests
- Invariant tests
- Concurrency tests
"""

import pytest
import asyncio
import json
import tempfile
import os
from pathlib import Path
from unittest.mock import Mock, MagicMock, AsyncMock, patch, call, ANY
from typing import Dict, List, Optional, Any
import yaml


# Mock the Rust module - in real scenario this would import from compiled module
# For test purposes, we create mock classes matching the contract


class Tool:
    """Mock MCP Tool structure"""
    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description


class ServerInfo:
    """Mock MCP ServerInfo structure"""
    def __init__(self, instructions: str, capabilities: Any):
        self.instructions = instructions
        self.capabilities = capabilities


class ServerCapabilities:
    """Mock MCP ServerCapabilities structure"""
    def __init__(self, tools: bool = False):
        self.tools = tools


class ListToolsResult:
    """Mock MCP ListToolsResult"""
    def __init__(self, tools: List[Tool], next_cursor: Optional[str] = None, meta: Optional[Dict] = None):
        self.tools = tools
        self.next_cursor = next_cursor
        self.meta = meta


class CallToolResult:
    """Mock MCP CallToolResult"""
    def __init__(self, content: List[Any]):
        self.content = content


class CallToolRequestParams:
    """Mock MCP CallToolRequestParams"""
    def __init__(self, name: str, arguments: Optional[Dict] = None):
        self.name = name
        self.arguments = arguments or {}


class PaginatedRequestParams:
    """Mock pagination parameters"""
    def __init__(self, cursor: Optional[str] = None):
        self.cursor = cursor


class RequestContext:
    """Mock MCP RequestContext"""
    def __init__(self, role: str = "server"):
        self.role = role


class McpError(Exception):
    """Mock MCP Error"""
    def __init__(self, error_type: str, message: str):
        self.error_type = error_type
        self.message = message
        super().__init__(f"{error_type}: {message}")


class UpstreamConfig:
    """Configuration for a single upstream MCP server process"""
    def __init__(self, command: str, args: List[str], env: Dict[str, str]):
        self.command = command
        self.args = args
        self.env = env


class ProxyConfig:
    """Configuration structure for the proxy server"""
    def __init__(self, servers: Dict[str, UpstreamConfig]):
        self.servers = servers


class UpstreamConnection:
    """Runtime state for a connected upstream MCP server"""
    def __init__(self, name: str, client: Any, tools: List[Tool]):
        self.name = name
        self.client = client
        self.tools = tools


class ProxyServer:
    """Main proxy server implementation"""
    def __init__(self, upstreams: Any, policy_path: Path, vault: Optional[Any]):
        self.upstreams = upstreams
        self.policy_path = policy_path
        self.vault = vault
    
    def get_info(self) -> ServerInfo:
        """Returns MCP server metadata"""
        return ServerInfo(
            instructions='Signet-proxied MCP server. All tool calls pass through policy enforcement.',
            capabilities=ServerCapabilities(tools=True)
        )
    
    async def list_tools(self, params: Optional[PaginatedRequestParams], context: RequestContext) -> ListToolsResult:
        """Aggregates tools from all upstream servers"""
        all_tools = []
        
        # Lock and iterate upstreams
        upstreams_list = await self.upstreams.get()
        
        for upstream in upstreams_list:
            for tool in upstream.tools:
                prefixed_tool = Tool(
                    name=f"{upstream.name}__{tool.name}",
                    description=f"[{upstream.name}] {tool.description}"
                )
                all_tools.append(prefixed_tool)
        
        return ListToolsResult(tools=all_tools, next_cursor=None, meta=None)
    
    async def call_tool(self, request: CallToolRequestParams, context: RequestContext) -> CallToolResult:
        """Intercepts tool calls and enforces policy"""
        # Validate tool name format
        if '__' not in request.name:
            raise McpError('invalid_tool_name_format', 'Tool name must contain __ separator')
        
        # Parse server and tool name
        parts = request.name.split('__', 1)
        server_name = parts[0]
        tool_name = parts[1]
        
        # Find upstream
        upstreams_list = await self.upstreams.get()
        upstream = None
        for u in upstreams_list:
            if u.name == server_name:
                upstream = u
                break
        
        if not upstream:
            raise McpError('unknown_server', f'Server {server_name} not found')
        
        # Load and check policy
        policy_decision = await self._check_policy(request)
        
        # Log to vault if exists
        if self.vault:
            detail = json.dumps(request.arguments)[:500]
            amount = 0.0 if policy_decision != 'Allow' else 0.0
            await self.vault.log_action(
                tool_name=request.name,
                decision=policy_decision,
                category='tool_call',
                amount=amount,
                detail=detail
            )
        
        # Handle policy decision
        if policy_decision != 'Allow':
            raise McpError('policy_denied', f'Policy decision: {policy_decision}')
        
        # Forward to upstream
        result = await upstream.client.call_tool(tool_name, request.arguments)
        return result
    
    async def _check_policy(self, request: CallToolRequestParams) -> str:
        """Check policy with integrity validation"""
        if not self.vault:
            return await default_policy(request)
        
        # Read policy file
        try:
            with open(self.policy_path, 'r') as f:
                policy_content = f.read()
            
            # Verify integrity
            if not await self.vault.verify_integrity(self.policy_path):
                return await default_policy(request)
            
            # Parse and evaluate policy
            return await evaluate_policy(policy_content, request)
        except Exception:
            return await default_policy(request)


async def default_policy(request: CallToolRequestParams) -> str:
    """Default policy when vault is None or integrity fails"""
    return 'Deny'


async def evaluate_policy(policy_content: str, request: CallToolRequestParams) -> str:
    """Evaluate policy content against request"""
    # Mock implementation
    return 'Allow'


def proxy_config_path() -> Path:
    """Returns the filesystem path to proxy.yaml"""
    return signet_dir().joinpath('proxy.yaml')


def signet_dir() -> Path:
    """Returns the Signet directory path"""
    return Path.home().joinpath('.signet')


async def connect_upstream(name: str, config: UpstreamConfig) -> UpstreamConnection:
    """Spawns an upstream MCP server and establishes connection"""
    # This would normally spawn a process and connect
    # Mock implementation for testing
    pass


async def run_proxy() -> None:
    """Main entry point for MCP proxy server"""
    # Load config
    config_path = proxy_config_path()
    
    if not config_path.exists():
        return
    
    try:
        with open(config_path, 'r') as f:
            config_data = yaml.safe_load(f)
        
        config = ProxyConfig(servers={
            name: UpstreamConfig(**server_config)
            for name, server_config in config_data.get('servers', {}).items()
        })
    except Exception as e:
        raise Exception(f'config_parse_error: {e}')
    
    if not config.servers:
        return
    
    # Connect upstreams
    upstreams = []
    for name, server_config in config.servers.items():
        try:
            upstream = await connect_upstream(name, server_config)
            upstreams.append(upstream)
        except Exception as e:
            # Log to stderr and continue
            print(f"Failed to connect to {name}: {e}", file=sys.stderr)
    
    # Create proxy server
    # ... implementation continues


class MockMutex:
    """Mock for Arc<Mutex<T>>"""
    def __init__(self, value):
        self._value = value
    
    async def get(self):
        return self._value


# ============================================================================
# TEST SUITE
# ============================================================================


class TestProxyConfigPath:
    """Tests for proxy_config_path function"""
    
    def test_proxy_config_path_happy(self):
        """Verify proxy_config_path returns correct path to proxy.yaml in Signet directory"""
        with patch('src_mcp_proxy.signet_dir') as mock_signet_dir:
            mock_signet_dir.return_value = Path('/home/user/.signet')
            
            result = proxy_config_path()
            
            assert result == Path('/home/user/.signet/proxy.yaml')
            assert result.name == 'proxy.yaml'
            assert result.parent == Path('/home/user/.signet')


class TestGetInfo:
    """Tests for get_info method"""
    
    def test_get_info_happy(self):
        """Verify get_info returns ServerInfo with correct instructions and capabilities"""
        upstreams = MockMutex([])
        policy_path = Path('/tmp/policy.yaml')
        vault = None
        
        proxy = ProxyServer(upstreams, policy_path, vault)
        result = proxy.get_info()
        
        assert result.instructions == 'Signet-proxied MCP server. All tool calls pass through policy enforcement.'
        assert result.capabilities.tools is True


class TestListTools:
    """Tests for list_tools method"""
    
    @pytest.mark.asyncio
    async def test_list_tools_single_upstream_happy(self):
        """Verify list_tools aggregates tools from single upstream with correct name/description formatting"""
        mock_client = Mock()
        tools = [
            Tool('read_file', 'Read a file from disk'),
            Tool('write_file', 'Write content to file')
        ]
        upstream = UpstreamConnection('filesystem', mock_client, tools)
        upstreams = MockMutex([upstream])
        
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        result = await proxy.list_tools(None, RequestContext())
        
        assert len(result.tools) == 2
        assert result.tools[0].name == 'filesystem__read_file'
        assert result.tools[1].name == 'filesystem__write_file'
        assert result.tools[0].description.startswith('[filesystem]')
        assert result.tools[1].description.startswith('[filesystem]')
        assert result.next_cursor is None
        assert result.meta is None
    
    @pytest.mark.asyncio
    async def test_list_tools_multiple_upstreams_happy(self):
        """Verify list_tools aggregates tools from multiple upstreams correctly"""
        upstream1 = UpstreamConnection('server1', Mock(), [
            Tool('tool1', 'Tool 1'),
            Tool('tool2', 'Tool 2')
        ])
        upstream2 = UpstreamConnection('server2', Mock(), [
            Tool('tool3', 'Tool 3')
        ])
        upstreams = MockMutex([upstream1, upstream2])
        
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        result = await proxy.list_tools(None, RequestContext())
        
        assert len(result.tools) == 3
        assert result.tools[0].name == 'server1__tool1'
        assert result.tools[1].name == 'server1__tool2'
        assert result.tools[2].name == 'server2__tool3'
        assert result.next_cursor is None
    
    @pytest.mark.asyncio
    async def test_list_tools_empty_upstreams(self):
        """Verify list_tools returns empty result when no upstreams are connected"""
        upstreams = MockMutex([])
        
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        result = await proxy.list_tools(None, RequestContext())
        
        assert len(result.tools) == 0
        assert result.next_cursor is None
        assert result.meta is None
    
    @pytest.mark.asyncio
    async def test_list_tools_upstream_with_no_tools(self):
        """Verify list_tools handles upstreams that expose zero tools"""
        upstream = UpstreamConnection('empty_server', Mock(), [])
        upstreams = MockMutex([upstream])
        
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        result = await proxy.list_tools(None, RequestContext())
        
        assert len(result.tools) == 0
    
    @pytest.mark.asyncio
    async def test_pagination_none_params(self):
        """Verify list_tools handles None pagination params correctly"""
        upstream = UpstreamConnection('server', Mock(), [
            Tool('t1', 'Tool 1'),
            Tool('t2', 'Tool 2')
        ])
        upstreams = MockMutex([upstream])
        
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        result = await proxy.list_tools(None, RequestContext())
        
        assert len(result.tools) == 2
        assert result.next_cursor is None
        assert result.meta is None


class TestCallTool:
    """Tests for call_tool method"""
    
    @pytest.mark.asyncio
    async def test_call_tool_valid_allow_decision_happy(self):
        """Verify call_tool forwards request to upstream when policy allows"""
        mock_client = AsyncMock()
        mock_client.call_tool = AsyncMock(return_value=CallToolResult([{'text': 'success'}]))
        
        upstream = UpstreamConnection('myserver', mock_client, [Tool('mytool', 'My tool')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Allow'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            # Create policy file
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request = CallToolRequestParams('myserver__mytool', {'arg': 'value'})
                result = await proxy.call_tool(request, RequestContext())
                
                assert mock_client.call_tool.called
                assert mock_vault.log_action.called
                
                # Verify vault logging
                log_call = mock_vault.log_action.call_args
                assert log_call[1]['tool_name'] == 'myserver__mytool'
                assert log_call[1]['decision'] == 'Allow'
                assert log_call[1]['amount'] == 0.0
                assert len(log_call[1]['detail']) <= 500
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_call_tool_invalid_format_error(self):
        """Verify call_tool returns invalid_tool_name_format error when name lacks '__' separator"""
        upstreams = MockMutex([])
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        
        request = CallToolRequestParams('invalid_tool_name', {})
        
        with pytest.raises(McpError) as exc_info:
            await proxy.call_tool(request, RequestContext())
        
        assert exc_info.value.error_type == 'invalid_tool_name_format'
    
    @pytest.mark.asyncio
    async def test_call_tool_unknown_server_error(self):
        """Verify call_tool returns unknown_server error when server prefix doesn't match any upstream"""
        upstream = UpstreamConnection('known_server', Mock(), [])
        upstreams = MockMutex([upstream])
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        
        request = CallToolRequestParams('unknown_server__tool', {})
        
        with pytest.raises(McpError) as exc_info:
            await proxy.call_tool(request, RequestContext())
        
        assert exc_info.value.error_type == 'unknown_server'
    
    @pytest.mark.asyncio
    async def test_call_tool_deny_decision(self):
        """Verify call_tool returns error for Deny policy decision and logs with amount 0.0"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Deny'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request = CallToolRequestParams('server__tool', {})
                
                with pytest.raises(McpError) as exc_info:
                    await proxy.call_tool(request, RequestContext())
                
                assert exc_info.value.error_type == 'policy_denied'
                assert not mock_client.call_tool.called
                
                # Verify vault logged with amount 0.0
                log_call = mock_vault.log_action.call_args
                assert log_call[1]['decision'] == 'Deny'
                assert log_call[1]['amount'] == 0.0
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_call_tool_ask_decision(self):
        """Verify call_tool returns error for Ask policy decision"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Ask'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request = CallToolRequestParams('server__tool', {})
                
                with pytest.raises(McpError):
                    await proxy.call_tool(request, RequestContext())
                
                assert not mock_client.call_tool.called
                
                log_call = mock_vault.log_action.call_args
                assert log_call[1]['amount'] == 0.0
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_call_tool_gate_decision(self):
        """Verify call_tool returns error for Gate policy decision"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Gate'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request = CallToolRequestParams('server__tool', {})
                
                with pytest.raises(McpError):
                    await proxy.call_tool(request, RequestContext())
                
                assert not mock_client.call_tool.called
                
                log_call = mock_vault.log_action.call_args
                assert log_call[1]['amount'] == 0.0
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_call_tool_ensure_decision(self):
        """Verify call_tool returns error for Ensure policy decision"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Ensure'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request = CallToolRequestParams('server__tool', {})
                
                with pytest.raises(McpError):
                    await proxy.call_tool(request, RequestContext())
                
                assert not mock_client.call_tool.called
                
                log_call = mock_vault.log_action.call_args
                assert log_call[1]['amount'] == 0.0
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_call_tool_policy_tampered(self):
        """Verify call_tool uses default_policy when integrity check fails"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=False)  # Integrity fails
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.default_policy', new_callable=AsyncMock) as mock_default:
            mock_default.return_value = 'Deny'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request = CallToolRequestParams('server__tool', {})
                
                with pytest.raises(McpError):
                    await proxy.call_tool(request, RequestContext())
                
                # Verify default_policy was called
                assert mock_default.called
                assert not mock_client.call_tool.called
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_call_tool_no_vault(self):
        """Verify call_tool uses default_policy and skips logging when vault is None"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        with patch('src_mcp_proxy.default_policy', new_callable=AsyncMock) as mock_default:
            mock_default.return_value = 'Deny'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)  # No vault
            
            request = CallToolRequestParams('server__tool', {})
            
            with pytest.raises(McpError):
                await proxy.call_tool(request, RequestContext())
            
            # Verify default_policy was used
            assert mock_default.called
            assert not mock_client.call_tool.called
    
    @pytest.mark.asyncio
    async def test_call_tool_detail_truncation(self):
        """Verify call_tool truncates detail to 500 characters in vault log"""
        mock_client = AsyncMock()
        mock_client.call_tool = AsyncMock(return_value=CallToolResult([]))
        
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Allow'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                # Create request with very long arguments
                long_arg = 'x' * 1000
                request = CallToolRequestParams('server__tool', {'data': long_arg})
                
                await proxy.call_tool(request, RequestContext())
                
                # Verify detail is truncated to 500 chars
                log_call = mock_vault.log_action.call_args
                detail = log_call[1]['detail']
                assert len(detail) == 500
            finally:
                os.unlink(policy_path)


class TestConnectUpstream:
    """Tests for connect_upstream function"""
    
    @pytest.mark.asyncio
    async def test_connect_upstream_happy(self):
        """Verify connect_upstream successfully spawns process and establishes connection"""
        config = UpstreamConfig(
            command='/usr/bin/python3',
            args=['-m', 'mcp_server'],
            env={'VAR': 'value'}
        )
        
        with patch('src_mcp_proxy.asyncio.create_subprocess_exec') as mock_subprocess, \
             patch('src_mcp_proxy.TokioChildProcess') as mock_transport, \
             patch('src_mcp_proxy.serve_client') as mock_serve:
            
            # Mock process
            mock_process = AsyncMock()
            mock_process.stdin = Mock()
            mock_process.stdout = Mock()
            mock_subprocess.return_value = mock_process
            
            # Mock client
            mock_client = AsyncMock()
            mock_client.list_tools = AsyncMock(return_value=ListToolsResult([
                Tool('tool1', 'Tool 1')
            ]))
            mock_serve.return_value = mock_client
            
            result = await connect_upstream('testserver', config)
            
            assert result.name == 'testserver'
            assert result.client is not None
            assert len(result.tools) == 1
            assert result.tools[0].name == 'tool1'
    
    @pytest.mark.asyncio
    async def test_connect_upstream_process_spawn_error(self):
        """Verify connect_upstream returns process_spawn_error for invalid command"""
        config = UpstreamConfig(
            command='/nonexistent/command',
            args=[],
            env={}
        )
        
        with patch('src_mcp_proxy.asyncio.create_subprocess_exec') as mock_subprocess:
            mock_subprocess.side_effect = FileNotFoundError('Command not found')
            
            with pytest.raises(Exception) as exc_info:
                await connect_upstream('testserver', config)
            
            assert 'process_spawn_error' in str(exc_info.value) or isinstance(exc_info.value, FileNotFoundError)
    
    @pytest.mark.asyncio
    async def test_connect_upstream_transport_error(self):
        """Verify connect_upstream returns transport_error when TokioChildProcess creation fails"""
        config = UpstreamConfig(command='/bin/echo', args=[], env={})
        
        with patch('src_mcp_proxy.asyncio.create_subprocess_exec') as mock_subprocess, \
             patch('src_mcp_proxy.TokioChildProcess') as mock_transport:
            
            mock_process = AsyncMock()
            mock_subprocess.return_value = mock_process
            
            mock_transport.side_effect = Exception('transport_error')
            
            with pytest.raises(Exception) as exc_info:
                await connect_upstream('testserver', config)
            
            assert 'transport_error' in str(exc_info.value)
    
    @pytest.mark.asyncio
    async def test_connect_upstream_client_connection_error(self):
        """Verify connect_upstream returns client_connection_error when serve_client fails"""
        config = UpstreamConfig(command='/bin/echo', args=[], env={})
        
        with patch('src_mcp_proxy.asyncio.create_subprocess_exec') as mock_subprocess, \
             patch('src_mcp_proxy.TokioChildProcess') as mock_transport, \
             patch('src_mcp_proxy.serve_client') as mock_serve:
            
            mock_process = AsyncMock()
            mock_subprocess.return_value = mock_process
            mock_transport.return_value = Mock()
            
            mock_serve.side_effect = Exception('client_connection_error')
            
            with pytest.raises(Exception) as exc_info:
                await connect_upstream('testserver', config)
            
            assert 'client_connection_error' in str(exc_info.value)
    
    @pytest.mark.asyncio
    async def test_connect_upstream_list_tools_error(self):
        """Verify connect_upstream returns list_tools_error when upstream doesn't respond to list_tools"""
        config = UpstreamConfig(command='/bin/echo', args=[], env={})
        
        with patch('src_mcp_proxy.asyncio.create_subprocess_exec') as mock_subprocess, \
             patch('src_mcp_proxy.TokioChildProcess') as mock_transport, \
             patch('src_mcp_proxy.serve_client') as mock_serve:
            
            mock_process = AsyncMock()
            mock_subprocess.return_value = mock_process
            mock_transport.return_value = Mock()
            
            mock_client = AsyncMock()
            mock_client.list_tools = AsyncMock(side_effect=Exception('list_tools_error'))
            mock_serve.return_value = mock_client
            
            with pytest.raises(Exception) as exc_info:
                await connect_upstream('testserver', config)
            
            assert 'list_tools_error' in str(exc_info.value)
    
    @pytest.mark.asyncio
    async def test_connect_upstream_env_variables(self):
        """Verify connect_upstream passes environment variables to child process"""
        config = UpstreamConfig(
            command='/bin/echo',
            args=['test'],
            env={'CUSTOM_VAR': 'custom_value', 'ANOTHER': 'value2'}
        )
        
        with patch('src_mcp_proxy.asyncio.create_subprocess_exec') as mock_subprocess, \
             patch('src_mcp_proxy.TokioChildProcess') as mock_transport, \
             patch('src_mcp_proxy.serve_client') as mock_serve:
            
            mock_process = AsyncMock()
            mock_subprocess.return_value = mock_process
            mock_transport.return_value = Mock()
            
            mock_client = AsyncMock()
            mock_client.list_tools = AsyncMock(return_value=ListToolsResult([]))
            mock_serve.return_value = mock_client
            
            await connect_upstream('testserver', config)
            
            # Verify environment variables were passed
            call_args = mock_subprocess.call_args
            assert 'env' in call_args[1] or len(call_args) > 0


class TestRunProxy:
    """Tests for run_proxy function"""
    
    @pytest.mark.asyncio
    async def test_run_proxy_no_config_file(self):
        """Verify run_proxy returns Ok(()) early when proxy.yaml doesn't exist"""
        with patch('src_mcp_proxy.proxy_config_path') as mock_path:
            mock_path.return_value = Path('/nonexistent/proxy.yaml')
            
            result = await run_proxy()
            
            assert result is None  # Returns early without error
    
    @pytest.mark.asyncio
    async def test_run_proxy_empty_servers(self):
        """Verify run_proxy returns Ok(()) early when config.servers is empty"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump({'servers': {}}, f)
            config_path = Path(f.name)
        
        try:
            with patch('src_mcp_proxy.proxy_config_path') as mock_path:
                mock_path.return_value = config_path
                
                result = await run_proxy()
                
                assert result is None  # Returns early
        finally:
            os.unlink(config_path)
    
    @pytest.mark.asyncio
    async def test_run_proxy_invalid_yaml(self):
        """Verify run_proxy returns config_parse_error for malformed YAML"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            f.write('invalid: yaml: content: [')
            config_path = Path(f.name)
        
        try:
            with patch('src_mcp_proxy.proxy_config_path') as mock_path:
                mock_path.return_value = config_path
                
                with pytest.raises(Exception) as exc_info:
                    await run_proxy()
                
                assert 'config_parse_error' in str(exc_info.value)
        finally:
            os.unlink(config_path)
    
    @pytest.mark.asyncio
    async def test_run_proxy_invalid_schema(self):
        """Verify run_proxy returns config_parse_error when YAML doesn't match ProxyConfig schema"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            # Missing required fields
            yaml.dump({'wrong_field': 'value'}, f)
            config_path = Path(f.name)
        
        try:
            with patch('src_mcp_proxy.proxy_config_path') as mock_path:
                mock_path.return_value = config_path
                
                with pytest.raises(Exception) as exc_info:
                    await run_proxy()
                
                assert 'config_parse_error' in str(exc_info.value) or 'servers' in str(exc_info.value)
        finally:
            os.unlink(config_path)
    
    @pytest.mark.asyncio
    async def test_run_proxy_partial_upstream_failure(self):
        """Verify run_proxy continues when some upstream connections fail"""
        config_data = {
            'servers': {
                'working_server': {
                    'command': '/bin/echo',
                    'args': [],
                    'env': {}
                },
                'failing_server': {
                    'command': '/nonexistent',
                    'args': [],
                    'env': {}
                },
                'another_working': {
                    'command': '/bin/true',
                    'args': [],
                    'env': {}
                }
            }
        }
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(config_data, f)
            config_path = Path(f.name)
        
        try:
            with patch('src_mcp_proxy.proxy_config_path') as mock_path, \
                 patch('src_mcp_proxy.connect_upstream') as mock_connect, \
                 patch('src_mcp_proxy.serve_server') as mock_serve, \
                 patch('sys.stderr'):
                
                mock_path.return_value = config_path
                
                # First and third succeed, second fails
                mock_connect.side_effect = [
                    UpstreamConnection('working_server', Mock(), []),
                    Exception('Connection failed'),
                    UpstreamConnection('another_working', Mock(), [])
                ]
                
                mock_service = AsyncMock()
                mock_service.waiting = AsyncMock(return_value=None)
                mock_serve.return_value = mock_service
                
                await run_proxy()
                
                # Verify all three connection attempts were made
                assert mock_connect.call_count == 3
        finally:
            os.unlink(config_path)
    
    @pytest.mark.asyncio
    async def test_run_proxy_happy(self):
        """Verify run_proxy successfully starts with valid config and upstreams"""
        config_data = {
            'servers': {
                'server1': {
                    'command': '/bin/echo',
                    'args': ['arg1'],
                    'env': {'VAR': 'val'}
                }
            }
        }
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(config_data, f)
            config_path = Path(f.name)
        
        try:
            with patch('src_mcp_proxy.proxy_config_path') as mock_path, \
                 patch('src_mcp_proxy.connect_upstream') as mock_connect, \
                 patch('src_mcp_proxy.serve_server') as mock_serve, \
                 patch('src_mcp_proxy.signet_dir') as mock_signet:
                
                mock_path.return_value = config_path
                mock_signet.return_value = Path('/tmp/.signet')
                
                mock_connect.return_value = UpstreamConnection('server1', Mock(), [])
                
                mock_service = AsyncMock()
                mock_service.waiting = AsyncMock(return_value=None)
                mock_serve.return_value = mock_service
                
                await run_proxy()
                
                # Verify upstream was connected
                assert mock_connect.called
                
                # Verify server was started
                assert mock_serve.called
                
                # Verify service.waiting was called (blocking)
                assert mock_service.waiting.called
        finally:
            os.unlink(config_path)


class TestConfigDeserialization:
    """Tests for ProxyConfig and UpstreamConfig deserialization"""
    
    @pytest.mark.asyncio
    async def test_config_deserialization_single_server(self):
        """Verify ProxyConfig deserializes single server configuration correctly"""
        config_data = {
            'servers': {
                'myserver': {
                    'command': '/usr/bin/python3',
                    'args': ['-m', 'server'],
                    'env': {'KEY': 'value'}
                }
            }
        }
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(config_data, f)
            config_path = Path(f.name)
        
        try:
            with open(config_path, 'r') as f:
                loaded_data = yaml.safe_load(f)
            
            assert 'servers' in loaded_data
            assert len(loaded_data['servers']) == 1
            assert 'myserver' in loaded_data['servers']
            
            server_config = loaded_data['servers']['myserver']
            assert server_config['command'] == '/usr/bin/python3'
            assert server_config['args'] == ['-m', 'server']
            assert server_config['env'] == {'KEY': 'value'}
        finally:
            os.unlink(config_path)
    
    @pytest.mark.asyncio
    async def test_config_deserialization_multiple_servers(self):
        """Verify ProxyConfig deserializes multiple server configurations"""
        config_data = {
            'servers': {
                'server1': {'command': 'cmd1', 'args': ['a1'], 'env': {}},
                'server2': {'command': 'cmd2', 'args': ['a2'], 'env': {}},
                'server3': {'command': 'cmd3', 'args': [], 'env': {'E': 'V'}}
            }
        }
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(config_data, f)
            config_path = Path(f.name)
        
        try:
            with open(config_path, 'r') as f:
                loaded_data = yaml.safe_load(f)
            
            assert len(loaded_data['servers']) == 3
            assert 'server1' in loaded_data['servers']
            assert 'server2' in loaded_data['servers']
            assert 'server3' in loaded_data['servers']
        finally:
            os.unlink(config_path)
    
    @pytest.mark.asyncio
    async def test_config_empty_args_array(self):
        """Verify UpstreamConfig handles empty args array"""
        config = UpstreamConfig(command='/bin/echo', args=[], env={})
        
        assert config.command == '/bin/echo'
        assert config.args == []
        assert isinstance(config.args, list)
    
    @pytest.mark.asyncio
    async def test_config_empty_env_map(self):
        """Verify UpstreamConfig handles empty env HashMap"""
        config = UpstreamConfig(command='/bin/echo', args=['test'], env={})
        
        assert config.command == '/bin/echo'
        assert config.env == {}
        assert isinstance(config.env, dict)


class TestInvariants:
    """Tests for contract invariants"""
    
    @pytest.mark.asyncio
    async def test_invariant_tool_name_format(self):
        """Verify all tools exposed by proxy always use server__tool format"""
        upstreams = [
            UpstreamConnection('srv1', Mock(), [Tool('t1', 'd1'), Tool('t2', 'd2')]),
            UpstreamConnection('srv2', Mock(), [Tool('t3', 'd3')])
        ]
        
        proxy = ProxyServer(MockMutex(upstreams), Path('/tmp/policy.yaml'), None)
        result = await proxy.list_tools(None, RequestContext())
        
        for tool in result.tools:
            assert '__' in tool.name
            parts = tool.name.split('__')
            assert len(parts) == 2
            assert parts[0] in ['srv1', 'srv2']
            assert parts[1] in ['t1', 't2', 't3']
    
    @pytest.mark.asyncio
    async def test_invariant_policy_reload(self):
        """Verify policy is reloaded and integrity-checked on every call_tool invocation"""
        mock_client = AsyncMock()
        mock_client.call_tool = AsyncMock(return_value=CallToolResult([]))
        
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Allow'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                # Make multiple calls
                for _ in range(3):
                    request = CallToolRequestParams('server__tool', {})
                    await proxy.call_tool(request, RequestContext())
                
                # Verify integrity check was called each time
                assert mock_vault.verify_integrity.call_count == 3
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_invariant_no_vault_uses_default_policy(self):
        """Verify default_policy is always used when vault is None"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        with patch('src_mcp_proxy.default_policy', new_callable=AsyncMock) as mock_default:
            mock_default.return_value = 'Deny'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
            
            # Make multiple calls
            for _ in range(3):
                request = CallToolRequestParams('server__tool', {})
                try:
                    await proxy.call_tool(request, RequestContext())
                except McpError:
                    pass
            
            # Verify default_policy was called each time
            assert mock_default.call_count == 3
    
    @pytest.mark.asyncio
    async def test_invariant_non_allow_amount_zero(self):
        """Verify amount logged is 0.0 for any non-Allow decision"""
        mock_client = AsyncMock()
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        decisions = ['Deny', 'Ask', 'Gate', 'Ensure']
        
        for decision in decisions:
            with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
                mock_policy.return_value = decision
                
                proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
                
                with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                    f.write('rules: []')
                    policy_path = Path(f.name)
                
                proxy.policy_path = policy_path
                
                try:
                    request = CallToolRequestParams('server__tool', {})
                    
                    try:
                        await proxy.call_tool(request, RequestContext())
                    except McpError:
                        pass
                    
                    # Verify amount is 0.0
                    log_call = mock_vault.log_action.call_args
                    assert log_call[1]['amount'] == 0.0
                    assert log_call[1]['decision'] == decision
                finally:
                    os.unlink(policy_path)


class TestConcurrency:
    """Tests for concurrent operations"""
    
    @pytest.mark.asyncio
    async def test_concurrent_list_tools(self):
        """Verify list_tools handles concurrent requests correctly with mutex"""
        upstream = UpstreamConnection('server', Mock(), [
            Tool('t1', 'd1'),
            Tool('t2', 'd2')
        ])
        upstreams = MockMutex([upstream])
        
        proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), None)
        
        # Run multiple concurrent list_tools calls
        results = await asyncio.gather(*[
            proxy.list_tools(None, RequestContext())
            for _ in range(5)
        ])
        
        # All should succeed with same results
        assert len(results) == 5
        for result in results:
            assert len(result.tools) == 2
            assert result.tools[0].name == 'server__t1'
    
    @pytest.mark.asyncio
    async def test_concurrent_call_tool_different_upstreams(self):
        """Verify concurrent call_tool requests to different upstreams work correctly"""
        mock_client1 = AsyncMock()
        mock_client1.call_tool = AsyncMock(return_value=CallToolResult([{'result': '1'}]))
        
        mock_client2 = AsyncMock()
        mock_client2.call_tool = AsyncMock(return_value=CallToolResult([{'result': '2'}]))
        
        upstream1 = UpstreamConnection('server1', mock_client1, [Tool('tool', 'desc')])
        upstream2 = UpstreamConnection('server2', mock_client2, [Tool('tool', 'desc')])
        
        upstreams = MockMutex([upstream1, upstream2])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Allow'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                request1 = CallToolRequestParams('server1__tool', {})
                request2 = CallToolRequestParams('server2__tool', {})
                
                results = await asyncio.gather(
                    proxy.call_tool(request1, RequestContext()),
                    proxy.call_tool(request2, RequestContext())
                )
                
                assert len(results) == 2
                assert mock_client1.call_tool.called
                assert mock_client2.call_tool.called
                assert mock_vault.log_action.call_count == 2
            finally:
                os.unlink(policy_path)
    
    @pytest.mark.asyncio
    async def test_concurrent_call_tool_same_upstream(self):
        """Verify concurrent call_tool requests to same upstream are handled safely"""
        mock_client = AsyncMock()
        call_count = {'count': 0}
        
        async def mock_call_tool(tool_name, args):
            call_count['count'] += 1
            await asyncio.sleep(0.01)  # Small delay to test concurrency
            return CallToolResult([{'call': call_count['count']}])
        
        mock_client.call_tool = mock_call_tool
        
        upstream = UpstreamConnection('server', mock_client, [Tool('tool', 'desc')])
        upstreams = MockMutex([upstream])
        
        mock_vault = AsyncMock()
        mock_vault.verify_integrity = AsyncMock(return_value=True)
        mock_vault.log_action = AsyncMock()
        
        with patch('src_mcp_proxy.evaluate_policy', new_callable=AsyncMock) as mock_policy:
            mock_policy.return_value = 'Allow'
            
            proxy = ProxyServer(upstreams, Path('/tmp/policy.yaml'), mock_vault)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
                f.write('rules: []')
                policy_path = Path(f.name)
            
            proxy.policy_path = policy_path
            
            try:
                requests = [CallToolRequestParams('server__tool', {}) for _ in range(3)]
                
                results = await asyncio.gather(*[
                    proxy.call_tool(req, RequestContext())
                    for req in requests
                ])
                
                assert len(results) == 3
                assert call_count['count'] == 3
                assert mock_vault.log_action.call_count == 3
            finally:
                os.unlink(policy_path)


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
