"""Test WebSocket tunnel client."""

import pytest
import asyncio
import json
from unittest.mock import Mock, patch, AsyncMock, MagicMock
from mcpbundles_proxy.tunnel import TunnelClient


@pytest.mark.asyncio
async def test_tunnel_client_init():
    """Test tunnel client initialization."""
    client = TunnelClient("test_token")
    assert client.token == "test_token"
    assert client.ws is None
    assert client.pending_requests == {}
    assert client.retry_count == 0
    assert client.running is True


@pytest.mark.asyncio
async def test_handle_request_localhost_validation():
    """Test that non-localhost requests are rejected."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    # Try to forward to non-localhost address
    request = {
        "request_id": "req_123",
        "method": "GET",
        "path": "/api/data",
        "target": "192.168.1.100:8000",  # Not localhost
        "headers": {},
        "body": None
    }
    
    await client.handle_request(request)
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data["type"] == "http_response"
    assert sent_data["status"] == 500
    assert "localhost" in sent_data["body"]


@pytest.mark.asyncio
async def test_handle_request_success():
    """Test successful request forwarding."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    request = {
        "request_id": "req_123",
        "method": "GET",
        "path": "/api/data",
        "target": "localhost:8000",
        "headers": {"Content-Type": "application/json"},
        "body": None
    }
    
    with patch('mcpbundles_proxy.tunnel.aiohttp.ClientSession') as mock_session_class:
        # Create proper async context manager
        mock_response = MagicMock()
        mock_response.status = 200
        mock_response.headers = {"Content-Type": "application/json"}
        mock_response.text = AsyncMock(return_value='{"result": "success"}')
        
        async_ctx_mgr = MagicMock()
        async_ctx_mgr.__aenter__ = AsyncMock(return_value=mock_response)
        async_ctx_mgr.__aexit__ = AsyncMock(return_value=None)
        
        mock_session = MagicMock()
        mock_session.request = MagicMock(return_value=async_ctx_mgr)
        
        session_ctx_mgr = MagicMock()
        session_ctx_mgr.__aenter__ = AsyncMock(return_value=mock_session)
        session_ctx_mgr.__aexit__ = AsyncMock(return_value=None)
        
        mock_session_class.return_value = session_ctx_mgr
        
        await client.handle_request(request)
        
        # Should send success response
        client.ws.send.assert_called_once()
        sent_data = json.loads(client.ws.send.call_args[0][0])
        assert sent_data["type"] == "http_response"
        assert sent_data["status"] == 200
        assert sent_data["request_id"] == "req_123"


@pytest.mark.asyncio
async def test_handle_request_timeout():
    """Test request timeout handling."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    request = {
        "request_id": "req_123",
        "method": "GET",
        "path": "/api/slow",
        "target": "localhost:8000",
        "headers": {},
        "body": None
    }
    
    with patch('mcpbundles_proxy.tunnel.aiohttp.ClientSession') as mock_session_class:
        mock_session = MagicMock()
        mock_session.request.side_effect = asyncio.TimeoutError()
        
        session_ctx_mgr = MagicMock()
        session_ctx_mgr.__aenter__ = AsyncMock(return_value=mock_session)
        session_ctx_mgr.__aexit__ = AsyncMock(return_value=None)
        
        mock_session_class.return_value = session_ctx_mgr
        
        await client.handle_request(request)
        
        # Should send timeout error
        client.ws.send.assert_called_once()
        sent_data = json.loads(client.ws.send.call_args[0][0])
        assert sent_data["status"] == 500
        assert "timeout" in sent_data["body"].lower()


@pytest.mark.asyncio
async def test_handle_request_connection_error():
    """Test connection error handling."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    request = {
        "request_id": "req_123",
        "method": "GET",
        "path": "/api/data",
        "target": "localhost:8000",
        "headers": {},
        "body": None
    }
    
    with patch('mcpbundles_proxy.tunnel.aiohttp.ClientSession') as mock_session_class:
        from aiohttp import ClientError
        mock_session = MagicMock()
        mock_session.request.side_effect = ClientError("Connection refused")
        
        session_ctx_mgr = MagicMock()
        session_ctx_mgr.__aenter__ = AsyncMock(return_value=mock_session)
        session_ctx_mgr.__aexit__ = AsyncMock(return_value=None)
        
        mock_session_class.return_value = session_ctx_mgr
        
        await client.handle_request(request)
        
        # Should send error response
        client.ws.send.assert_called_once()
        sent_data = json.loads(client.ws.send.call_args[0][0])
        assert sent_data["status"] == 500
        assert "Connection failed" in sent_data["body"]


@pytest.mark.asyncio
async def test_handle_request_size_limit():
    """Test response size limit enforcement."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    request = {
        "request_id": "req_123",
        "method": "GET",
        "path": "/api/large",
        "target": "localhost:8000",
        "headers": {},
        "body": None
    }
    
    with patch('mcpbundles_proxy.tunnel.aiohttp.ClientSession') as mock_session_class:
        mock_response = MagicMock()
        mock_response.status = 200
        mock_response.headers = {}
        mock_response.text = AsyncMock(return_value="x" * (11 * 1024 * 1024))  # 11MB
        
        async_ctx_mgr = MagicMock()
        async_ctx_mgr.__aenter__ = AsyncMock(return_value=mock_response)
        async_ctx_mgr.__aexit__ = AsyncMock(return_value=None)
        
        mock_session = MagicMock()
        mock_session.request = MagicMock(return_value=async_ctx_mgr)
        
        session_ctx_mgr = MagicMock()
        session_ctx_mgr.__aenter__ = AsyncMock(return_value=mock_session)
        session_ctx_mgr.__aexit__ = AsyncMock(return_value=None)
        
        mock_session_class.return_value = session_ctx_mgr
        
        await client.handle_request(request)
        
        # Should send size limit error
        client.ws.send.assert_called_once()
        sent_data = json.loads(client.ws.send.call_args[0][0])
        assert sent_data["status"] == 500
        assert "10MB" in sent_data["body"]


@pytest.mark.asyncio
async def test_handle_disconnect():
    """Test cleanup on disconnect."""
    client = TunnelClient("test_token")
    
    # Add some pending requests
    future1 = asyncio.Future()
    future2 = asyncio.Future()
    client.pending_requests = {
        "req_1": future1,
        "req_2": future2
    }
    
    await client._handle_disconnect()
    
    # All pending requests should be failed
    assert len(client.pending_requests) == 0
    assert future1.done()
    assert future2.done()
    with pytest.raises(Exception, match="Tunnel disconnected"):
        future1.result()


@pytest.mark.asyncio
async def test_stop():
    """Test graceful stop."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    client.running = True
    
    await client.stop()
    
    assert client.running is False
    client.ws.close.assert_called_once()


@pytest.mark.asyncio
async def test_send_error_with_ws():
    """Test sending error when websocket is connected."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    await client.send_error("req_123", "Test error")
    
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data["type"] == "http_response"
    assert sent_data["request_id"] == "req_123"
    assert sent_data["status"] == 500
    assert "Test error" in sent_data["body"]


@pytest.mark.asyncio
async def test_send_error_without_ws():
    """Test sending error when websocket is None."""
    client = TunnelClient("test_token")
    client.ws = None
    
    # Should not raise exception
    await client.send_error("req_123", "Test error")


@pytest.mark.asyncio
async def test_handle_service_discovery():
    """Test handling service discovery request."""
    mock_service_manager = AsyncMock()
    mock_service_manager.discover_services = AsyncMock(return_value={
        '5432': {'available': True, 'type': 'postgresql'},
        '6379': {'available': False, 'type': None}
    })
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_discovery({'ports': [5432, 6379]})
    
    # Should call service manager
    mock_service_manager.discover_services.assert_called_once_with([5432, 6379])
    
    # Should send response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'discovery_result'
    assert '5432' in sent_data['ports']


@pytest.mark.asyncio
async def test_handle_service_discovery_no_manager():
    """Test handling service discovery without service manager."""
    client = TunnelClient("test_token", service_manager=None)
    client.ws = AsyncMock()
    
    await client.handle_service_discovery({'ports': [5432]})
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert 'error' in sent_data


@pytest.mark.asyncio
async def test_handle_service_verify():
    """Test handling service verification request."""
    mock_service_manager = AsyncMock()
    mock_service_manager.verify_service = AsyncMock(return_value={
        'verified': True,
        'target': 'localhost:5432',
        'service_type': 'postgresql'
    })
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_verify({'target': 'localhost:5432'})
    
    # Should call service manager
    mock_service_manager.verify_service.assert_called_once_with('localhost:5432')
    
    # Should send response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'verify_result'
    assert sent_data['verified'] is True
    assert sent_data['target'] == 'localhost:5432'


@pytest.mark.asyncio
async def test_handle_service_start():
    """Test handling service start request."""
    mock_service_manager = AsyncMock()
    mock_service_manager.start_service = AsyncMock(return_value={
        'started': True,
        'service': 'browser',
        'port': 9223
    })
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_start({
        'service': 'browser',
        'config': {'headless': False}
    })
    
    # Should call service manager
    mock_service_manager.start_service.assert_called_once_with('browser', {'headless': False})
    
    # Should send response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'service_started'
    assert sent_data['started'] is True


@pytest.mark.asyncio
async def test_handle_service_stop():
    """Test handling service stop request."""
    mock_service_manager = AsyncMock()
    mock_service_manager.stop_service = AsyncMock(return_value={
        'stopped': True,
        'service': 'browser'
    })
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_stop({'service': 'browser'})
    
    # Should call service manager
    mock_service_manager.stop_service.assert_called_once_with('browser')
    
    # Should send response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'service_stopped'
    assert sent_data['stopped'] is True


@pytest.mark.asyncio
async def test_handle_service_update():
    """Test handling service update request."""
    mock_service_manager = AsyncMock()
    mock_service_manager.update_service = AsyncMock(return_value={
        'updated': True,
        'service': 'browser',
        'mode': 'hidden'
    })
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_update({
        'service': 'browser',
        'config': {'headless': True}
    })
    
    # Should call service manager
    mock_service_manager.update_service.assert_called_once_with('browser', {'headless': True})
    
    # Should send response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'service_updated'
    assert sent_data['updated'] is True


@pytest.mark.asyncio
async def test_send_status():
    """Test sending status update."""
    mock_service_manager = Mock()
    mock_service_manager.get_status.return_value = {
        'browser': {'enabled': True, 'mode': 'visible'}
    }
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.send_status()
    
    # Should send status update
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'status_update'
    assert 'browser' in sent_data['services']


@pytest.mark.asyncio
async def test_send_status_no_ws():
    """Test sending status without websocket."""
    client = TunnelClient("test_token")
    client.ws = None
    
    # Should not raise exception
    await client.send_status()


@pytest.mark.asyncio
async def test_send_response():
    """Test sending generic response."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    await client.send_response('test_response', {'key': 'value'})
    
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'test_response'
    assert sent_data['key'] == 'value'


@pytest.mark.asyncio
async def test_listen_http_request():
    """Test listen handles http_request messages."""
    client = TunnelClient("test_token")
    
    messages = [
        json.dumps({"type": "http_request", "request_id": "123", "method": "GET", "path": "/", "target": "localhost:8000"})
    ]
    
    async def async_iter_messages():
        for msg in messages:
            yield msg
    
    mock_ws = AsyncMock()
    mock_ws.__aiter__ = lambda self: async_iter_messages()
    client.ws = mock_ws
    
    # Mock handle_request
    client.handle_request = AsyncMock()
    
    # Listen should process one message then stop when iterator ends
    await client.listen()
    
    client.handle_request.assert_called_once()


@pytest.mark.asyncio
async def test_listen_invalid_json():
    """Test listen handles invalid JSON gracefully."""
    client = TunnelClient("test_token")
    
    messages = [
        "invalid json {{{",
        json.dumps({"type": "status_request"})
    ]
    
    async def async_iter_messages():
        for msg in messages:
            yield msg
    
    mock_ws = AsyncMock()
    mock_ws.__aiter__ = lambda self: async_iter_messages()
    client.ws = mock_ws
    
    client.send_status = AsyncMock()
    
    await client.listen()
    
    # Should handle invalid JSON and continue to process next message
    client.send_status.assert_called_once()


@pytest.mark.asyncio
async def test_listen_unknown_message_type():
    """Test listen handles unknown message types."""
    client = TunnelClient("test_token")
    
    messages = [
        json.dumps({"type": "unknown_type", "data": "test"})
    ]
    
    async def async_iter_messages():
        for msg in messages:
            yield msg
    
    mock_ws = AsyncMock()
    mock_ws.__aiter__ = lambda self: async_iter_messages()
    client.ws = mock_ws
    
    # Should not crash, just log warning
    await client.listen()


@pytest.mark.asyncio
async def test_handle_config_update():
    """Test handling config update message."""
    mock_service_manager = AsyncMock()
    mock_service_manager.update_config = AsyncMock()
    mock_service_manager.get_status = Mock(return_value={"browser": {"enabled": True}})
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_config_update({"config": {"browser": {"enabled": True}}})
    
    mock_service_manager.update_config.assert_called_once_with({"browser": {"enabled": True}})
    client.ws.send.assert_called_once()


@pytest.mark.asyncio
async def test_handle_config_update_no_manager():
    """Test config update without service manager."""
    client = TunnelClient("test_token", service_manager=None)
    client.ws = AsyncMock()
    
    # Should not crash
    await client.handle_config_update({"config": {}})


@pytest.mark.asyncio
async def test_handle_config_update_error():
    """Test config update with error."""
    mock_service_manager = AsyncMock()
    mock_service_manager.update_config = AsyncMock(side_effect=Exception("Config error"))
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    # Should not crash
    await client.handle_config_update({"config": {}})


@pytest.mark.asyncio
async def test_send_status_no_service_manager():
    """Test sending status without service manager."""
    client = TunnelClient("test_token", service_manager=None)
    client.ws = AsyncMock()
    
    await client.send_status()
    
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['type'] == 'status_update'
    assert sent_data['services'] == {}


@pytest.mark.asyncio
async def test_send_status_error():
    """Test error handling in send_status."""
    mock_service_manager = Mock()
    mock_service_manager.get_status.side_effect = Exception("Status error")
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    # Should not crash
    await client.send_status()


@pytest.mark.asyncio
async def test_send_response_no_ws():
    """Test send_response without websocket."""
    client = TunnelClient("test_token")
    client.ws = None
    
    # Should not crash
    await client.send_response('test', {})


@pytest.mark.asyncio
async def test_send_response_error():
    """Test send_response with send error."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    client.ws.send.side_effect = Exception("Send error")
    
    # Should not crash
    await client.send_response('test', {})


@pytest.mark.asyncio
async def test_handle_service_discovery_error():
    """Test service discovery with error."""
    mock_service_manager = AsyncMock()
    mock_service_manager.discover_services = AsyncMock(side_effect=Exception("Discovery error"))
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_discovery({'ports': [5432]})
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert 'error' in sent_data


@pytest.mark.asyncio
async def test_handle_service_verify_error():
    """Test service verify with error."""
    mock_service_manager = AsyncMock()
    mock_service_manager.verify_service = AsyncMock(side_effect=Exception("Verify error"))
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_verify({'target': 'localhost:5432'})
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['verified'] is False
    assert 'error' in sent_data


@pytest.mark.asyncio
async def test_handle_service_start_error():
    """Test service start with error."""
    mock_service_manager = AsyncMock()
    mock_service_manager.start_service = AsyncMock(side_effect=Exception("Start error"))
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_start({'service': 'browser', 'config': {}})
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['started'] is False
    assert 'error' in sent_data


@pytest.mark.asyncio
async def test_handle_service_start_no_manager():
    """Test service start without manager."""
    client = TunnelClient("test_token", service_manager=None)
    client.ws = AsyncMock()
    
    await client.handle_service_start({'service': 'browser', 'config': {}})
    
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['started'] is False
    assert 'Service manager not available' in sent_data['error']


@pytest.mark.asyncio
async def test_handle_service_stop_error():
    """Test service stop with error."""
    mock_service_manager = AsyncMock()
    mock_service_manager.stop_service = AsyncMock(side_effect=Exception("Stop error"))
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_stop({'service': 'browser'})
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['stopped'] is False
    assert 'error' in sent_data


@pytest.mark.asyncio
async def test_handle_service_stop_no_manager():
    """Test service stop without manager."""
    client = TunnelClient("test_token", service_manager=None)
    client.ws = AsyncMock()
    
    await client.handle_service_stop({'service': 'browser'})
    
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['stopped'] is False


@pytest.mark.asyncio
async def test_handle_service_update_error():
    """Test service update with error."""
    mock_service_manager = AsyncMock()
    mock_service_manager.update_service = AsyncMock(side_effect=Exception("Update error"))
    
    client = TunnelClient("test_token", service_manager=mock_service_manager)
    client.ws = AsyncMock()
    
    await client.handle_service_update({'service': 'browser', 'config': {}})
    
    # Should send error response
    client.ws.send.assert_called_once()
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['updated'] is False
    assert 'error' in sent_data


@pytest.mark.asyncio
async def test_handle_service_update_no_manager():
    """Test service update without manager."""
    client = TunnelClient("test_token", service_manager=None)
    client.ws = AsyncMock()
    
    await client.handle_service_update({'service': 'browser', 'config': {}})
    
    sent_data = json.loads(client.ws.send.call_args[0][0])
    assert sent_data['updated'] is False


@pytest.mark.asyncio
async def test_handle_request_general_exception():
    """Test handle_request with unexpected exception."""
    client = TunnelClient("test_token")
    client.ws = AsyncMock()
    
    request = {
        "request_id": "req_123",
        "method": "GET",
        "path": "/api/data",
        "target": "localhost:8000",
        "headers": {},
        "body": None
    }
    
    with patch('mcpbundles_proxy.tunnel.aiohttp.ClientSession') as mock_session_class:
        mock_session_class.side_effect = Exception("Unexpected error")
        
        await client.handle_request(request)
        
        # Should send error response
        client.ws.send.assert_called_once()
        sent_data = json.loads(client.ws.send.call_args[0][0])
        assert sent_data["status"] == 500
        assert "Internal error" in sent_data["body"]

