"""Tests for Tool Bridge Server."""

import pytest
import socket
import json
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, Mock, patch
from mcp_forge.mcp.bridge import ToolBridgeServer


@pytest.fixture
def temp_socket_path():
    """Create temporary socket path."""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield Path(tmpdir) / "test_bridge.sock"


@pytest.fixture
def mock_client_manager():
    """Create mock MCP client manager."""
    manager = AsyncMock()
    manager.call_tool = AsyncMock(return_value={"result": "success"})
    return manager


@pytest.fixture
def mock_audit_logger():
    """Create mock audit logger."""
    logger = Mock()
    logger.log = Mock()
    return logger


@pytest.mark.asyncio
async def test_bridge_server_starts_and_stops(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test that bridge server starts and stops cleanly."""
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    # Start server
    bridge.start()
    
    # Verify socket was created
    assert temp_socket_path.exists()
    
    # Stop server
    bridge.stop()
    
    # Verify socket was removed
    assert not temp_socket_path.exists()


@pytest.mark.asyncio
async def test_receive_and_forward_tool_call(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test receiving tool call request and forwarding to client."""
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    
    try:
        # Connect as client and send tool call request
        client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client_sock.connect(str(temp_socket_path))
        
        request = {
            "tool": "test_tool",
            "params": {"arg1": "value1", "arg2": 42}
        }
        
        client_sock.sendall(json.dumps(request).encode('utf-8'))
        client_sock.shutdown(socket.SHUT_WR)
        
        # Receive response
        response_data = b''
        while True:
            chunk = client_sock.recv(4096)
            if not chunk:
                break
            response_data += chunk
        
        response = json.loads(response_data.decode('utf-8'))
        
        # Verify response
        assert response["success"] is True
        assert response["result"] == {"result": "success"}
        
        # Verify tool was called with correct arguments
        mock_client_manager.call_tool.assert_called_once_with(
            "test_tool",
            {"arg1": "value1", "arg2": 42}
        )
        
        # Verify audit log was called
        mock_audit_logger.log.assert_called()
        
        client_sock.close()
    finally:
        bridge.stop()


@pytest.mark.asyncio
async def test_handle_tool_call_error(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test handling tool call errors."""
    # Make client manager raise error
    mock_client_manager.call_tool = AsyncMock(side_effect=RuntimeError("Tool failed"))
    
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    
    try:
        # Connect and send request
        client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client_sock.connect(str(temp_socket_path))
        
        request = {"tool": "failing_tool", "params": {}}
        client_sock.sendall(json.dumps(request).encode('utf-8'))
        client_sock.shutdown(socket.SHUT_WR)
        
        # Receive response
        response_data = b''
        while True:
            chunk = client_sock.recv(4096)
            if not chunk:
                break
            response_data += chunk
        
        response = json.loads(response_data.decode('utf-8'))
        
        # Verify error response
        assert response["success"] is False
        assert "Tool failed" in response["error"]
        
        client_sock.close()
    finally:
        bridge.stop()


@pytest.mark.asyncio
async def test_handle_invalid_json(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test handling invalid JSON in request."""
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    
    try:
        # Connect and send invalid JSON
        client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client_sock.connect(str(temp_socket_path))
        
        client_sock.sendall(b"not valid json")
        client_sock.shutdown(socket.SHUT_WR)
        
        # Receive response
        response_data = b''
        while True:
            chunk = client_sock.recv(4096)
            if not chunk:
                break
            response_data += chunk
        
        response = json.loads(response_data.decode('utf-8'))
        
        # Verify error response
        assert response["success"] is False
        assert "Invalid JSON" in response["error"]
        
        client_sock.close()
    finally:
        bridge.stop()


@pytest.mark.asyncio
async def test_handle_missing_tool_field(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test handling request missing 'tool' field."""
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    
    try:
        # Connect and send request without 'tool' field
        client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client_sock.connect(str(temp_socket_path))
        
        request = {"params": {"arg1": "value1"}}  # Missing 'tool'
        client_sock.sendall(json.dumps(request).encode('utf-8'))
        client_sock.shutdown(socket.SHUT_WR)
        
        # Receive response
        response_data = b''
        while True:
            chunk = client_sock.recv(4096)
            if not chunk:
                break
            response_data += chunk
        
        response = json.loads(response_data.decode('utf-8'))
        
        # Verify error response
        assert response["success"] is False
        assert "tool" in response["error"].lower()
        
        client_sock.close()
    finally:
        bridge.stop()


@pytest.mark.asyncio
async def test_concurrent_requests(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test handling multiple concurrent requests."""
    import threading
    
    # Track call counts
    call_count = {"count": 0}
    
    async def mock_call_tool(tool_name, arguments):
        call_count["count"] += 1
        return {"result": f"success_{call_count['count']}"}
    
    mock_client_manager.call_tool = mock_call_tool
    
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    
    try:
        results = []
        
        def make_request(tool_name):
            client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            client_sock.connect(str(temp_socket_path))
            
            request = {"tool": tool_name, "params": {}}
            client_sock.sendall(json.dumps(request).encode('utf-8'))
            client_sock.shutdown(socket.SHUT_WR)
            
            response_data = b''
            while True:
                chunk = client_sock.recv(4096)
                if not chunk:
                    break
                response_data += chunk
            
            response = json.loads(response_data.decode('utf-8'))
            results.append(response)
            client_sock.close()
        
        # Make 3 concurrent requests
        threads = []
        for i in range(3):
            thread = threading.Thread(target=make_request, args=(f"tool_{i}",))
            threads.append(thread)
            thread.start()
        
        # Wait for all threads
        for thread in threads:
            thread.join()
        
        # Verify all requests succeeded
        assert len(results) == 3
        for result in results:
            assert result["success"] is True
        
        # Verify all were processed
        assert call_count["count"] == 3
    finally:
        bridge.stop()


@pytest.mark.asyncio
async def test_audit_logging_tool_name_only(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test that audit log only logs tool name, not parameters."""
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    
    try:
        # Send request with sensitive parameters
        client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client_sock.connect(str(temp_socket_path))
        
        request = {
            "tool": "sensitive_tool",
            "params": {"password": "secret123", "token": "abc123"}
        }
        
        client_sock.sendall(json.dumps(request).encode('utf-8'))
        client_sock.shutdown(socket.SHUT_WR)
        
        # Receive response
        response_data = b''
        while True:
            chunk = client_sock.recv(4096)
            if not chunk:
                break
            response_data += chunk
        
        client_sock.close()
        
        # Verify audit log was called
        mock_audit_logger.log.assert_called()
        
        # Get the log call arguments
        log_call = mock_audit_logger.log.call_args
        
        # Verify tool name is in log
        log_str = str(log_call)
        assert "sensitive_tool" in log_str
        
        # Verify sensitive parameters are NOT in log
        assert "secret123" not in log_str
        assert "abc123" not in log_str
    finally:
        bridge.stop()


@pytest.mark.asyncio
async def test_socket_cleanup_on_error(temp_socket_path, mock_client_manager, mock_audit_logger):
    """Test that socket is cleaned up even if server encounters error."""
    bridge = ToolBridgeServer(
        socket_path=temp_socket_path,
        client_manager=mock_client_manager,
        audit_logger=mock_audit_logger
    )
    
    bridge.start()
    assert temp_socket_path.exists()
    
    # Stop should cleanup
    bridge.stop()
    assert not temp_socket_path.exists()