"""Tests for Execute Python Tool."""

import pytest
from unittest.mock import Mock, AsyncMock, patch
from mcp.types import Tool, TextContent
import json

from mcp_forge.server.tools.execute_python import ExecutePythonTool
from mcp_forge.execution.simple.backend import ExecutionResult


@pytest.fixture
def mock_simple_backend():
    """Mock simple backend."""
    backend = Mock()
    backend.execute = AsyncMock(return_value=ExecutionResult(
        success=True,
        stdout="Hello World",
        stderr="",
        result="42",
        execution_time=0.5,
        exit_code=0
    ))
    return backend


@pytest.fixture
def mock_jupyter_backend():
    """Mock jupyter backend."""
    backend = Mock()
    backend.execute = AsyncMock(return_value=ExecutionResult(
        success=True,
        stdout="Stateful execution",
        stderr="",
        result="session-123",
        execution_time=1.2,
        exit_code=0
    ))
    backend.execute_in_session = AsyncMock(return_value=ExecutionResult(
        success=True,
        stdout="Using existing session",
        stderr="",
        result="[1, 2, 3]",
        execution_time=0.3,
        exit_code=0
    ))
    return backend


@pytest.fixture
def mock_client_manager():
    """Mock MCP client manager."""
    manager = Mock()
    manager.list_all_tools = AsyncMock(return_value=[
        {"name": "github_search_repos", "description": "Search GitHub repositories"},
        {"name": "filesystem_read", "description": "Read file contents"}
    ])
    return manager


@pytest.fixture
def mock_bridge_server():
    """Mock tool bridge server."""
    server = Mock()
    server.socket_path = "/tmp/mcp-bridge.sock"
    server.is_running = Mock(return_value=True)
    return server


@pytest.fixture
def mock_injection_generator():
    """Mock tool injection generator."""
    generator = Mock()
    generator.generate_injection_code = Mock(return_value="""
# MCP Tool Injection
def github_search_repos(**kwargs):
    import socket
    # ... tool implementation
    pass
""")
    return generator


@pytest.fixture
def mock_config():
    """Mock forge configuration."""
    config = Mock()
    config.execution = Mock()
    config.execution.default_backend = "simple"
    config.execution.default_timeout = 300
    config.execution.max_timeout = 1800
    config.execution.default_memory = "512m"
    config.execution.max_memory = "2g"
    return config


@pytest.fixture
def execute_python_tool(
    mock_simple_backend,
    mock_jupyter_backend,
    mock_client_manager,
    mock_bridge_server,
    mock_injection_generator,
    mock_config
):
    """Create ExecutePythonTool instance with mocked dependencies."""
    return ExecutePythonTool(
        simple_backend=mock_simple_backend,
        jupyter_backend=mock_jupyter_backend,
        client_manager=mock_client_manager,
        bridge_server=mock_bridge_server,
        injection_generator=mock_injection_generator,
        config=mock_config
    )


@pytest.mark.asyncio
async def test_get_tool_definition(execute_python_tool):
    """Test that tool definition matches MCP spec."""
    definition = execute_python_tool.get_tool_definition()
    
    assert isinstance(definition, Tool)
    assert definition.name == "execute_python"
    assert definition.description is not None
    assert "Execute Python code" in definition.description
    
    # Verify required schema properties
    schema = definition.inputSchema
    assert schema["type"] == "object"
    assert "code" in schema["properties"]
    assert "mcp_tools" in schema["properties"]
    assert "session_id" in schema["properties"]
    assert "backend" in schema["properties"]
    assert "timeout" in schema["properties"]
    assert "custom_image" in schema["properties"]
    assert "environment" in schema["properties"]
    assert schema["required"] == ["code"]


@pytest.mark.asyncio
async def test_execute_simple_backend_stateless(
    execute_python_tool,
    mock_simple_backend
):
    """Test execution with simple backend (stateless)."""
    arguments = {
        "code": "print('Hello World'); 42"
    }
    
    result = await execute_python_tool.execute(arguments)
    
    assert isinstance(result, list)
    assert len(result) == 1
    assert isinstance(result[0], TextContent)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    assert response["stdout"] == "Hello World"
    assert response["result"] == "42"
    assert response["execution_time"] == 0.5
    assert "session_id" not in response or response["session_id"] is None
    
    # Verify simple backend was called
    mock_simple_backend.execute.assert_called_once()


@pytest.mark.asyncio
async def test_execute_jupyter_backend_stateful(
    execute_python_tool,
    mock_jupyter_backend
):
    """Test execution with jupyter backend (stateful)."""
    arguments = {
        "code": "x = 42; print('Stateful')",
        "session_id": "test-session",
        "backend": "jupyter"
    }
    
    result = await execute_python_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    assert "session_id" in response
    
    # Verify jupyter backend was called
    mock_jupyter_backend.execute_in_session.assert_called_once()


@pytest.mark.asyncio
async def test_execute_with_mcp_tools(
    execute_python_tool,
    mock_simple_backend,
    mock_injection_generator,
    mock_bridge_server
):
    """Test execution with MCP tool injection."""
    arguments = {
        "code": "repos = github_search_repos(query='test', max_results=10)",
        "mcp_tools": ["github_search_repos"]
    }
    
    result = await execute_python_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    assert "available_tools" in response
    assert "github_search_repos" in response["available_tools"]
    
    # Verify injection generator was called
    mock_injection_generator.generate_injection_code.assert_called_once_with(
        tool_names=["github_search_repos"],
        socket_path="/tmp/mcp-bridge.sock"
    )


@pytest.mark.asyncio
async def test_execute_with_custom_image(
    execute_python_tool,
    mock_simple_backend
):
    """Test execution with custom image."""
    arguments = {
        "code": "import pandas as pd; pd.DataFrame()",
        "custom_image": "mcp-forge/custom:my-ml-env"
    }
    
    _result = await execute_python_tool.execute(arguments)
    
    # Verify backend was called with custom image
    call_args = mock_simple_backend.execute.call_args
    assert call_args is not None
    # Check that custom_image was passed in execution options
    assert "image" in call_args.kwargs or "custom_image" in call_args.kwargs


@pytest.mark.asyncio
async def test_execute_with_environment_template(
    execute_python_tool,
    mock_simple_backend
):
    """Test execution with environment template."""
    arguments = {
        "code": "import numpy as np; np.array([1,2,3])",
        "environment": "datascience"
    }
    
    _result = await execute_python_tool.execute(arguments)
    
    # Verify backend was called with environment specification
    mock_simple_backend.execute.assert_called_once()
    call_args = mock_simple_backend.execute.call_args
    assert call_args is not None


@pytest.mark.asyncio
async def test_execute_with_timeout(
    execute_python_tool,
    mock_simple_backend
):
    """Test execution with custom timeout."""
    arguments = {
        "code": "import time; time.sleep(10)",
        "timeout": 5
    }
    
    _result = await execute_python_tool.execute(arguments)
    
    # Verify timeout was passed to backend
    call_args = mock_simple_backend.execute.call_args
    assert call_args is not None


@pytest.mark.asyncio
async def test_validate_arguments_missing_code(execute_python_tool):
    """Test that validation fails when code is missing."""
    arguments = {}
    
    with pytest.raises((ValueError, KeyError)):
        await execute_python_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_invalid_backend(execute_python_tool):
    """Test that validation fails for invalid backend."""
    arguments = {
        "code": "print('test')",
        "backend": "invalid"
    }
    
    with pytest.raises(ValueError, match="backend"):
        await execute_python_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_invalid_timeout(
    execute_python_tool,
    mock_config
):
    """Test that validation fails for timeout exceeding max."""
    arguments = {
        "code": "print('test')",
        "timeout": 3600  # Exceeds max_timeout of 1800
    }
    
    with pytest.raises(ValueError, match="Timeout"):
        await execute_python_tool.execute(arguments)


@pytest.mark.asyncio
async def test_backend_selection_with_session_id(execute_python_tool):
    """Test that jupyter backend is selected when session_id provided."""
    backend = execute_python_tool._select_backend(
        session_id="test-session",
        backend=None
    )
    
    assert backend == "jupyter"


@pytest.mark.asyncio
async def test_backend_selection_default(
    execute_python_tool,
    mock_config
):
    """Test that default backend is used when no session_id."""
    backend = execute_python_tool._select_backend(
        session_id=None,
        backend=None
    )
    
    assert backend == "simple"


@pytest.mark.asyncio
async def test_backend_selection_explicit(execute_python_tool):
    """Test that explicit backend is respected."""
    backend = execute_python_tool._select_backend(
        session_id=None,
        backend="jupyter"
    )
    
    assert backend == "jupyter"


@pytest.mark.asyncio
async def test_execution_error_handling(
    execute_python_tool,
    mock_simple_backend
):
    """Test that execution errors are handled gracefully."""
    mock_simple_backend.execute = AsyncMock(return_value=ExecutionResult(
        success=False,
        stdout="",
        stderr="NameError: name 'undefined_variable' is not defined",
        result=None,
        execution_time=0.1,
        exit_code=1
    ))
    
    arguments = {
        "code": "print(undefined_variable)"
    }
    
    result = await execute_python_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is False
    assert "NameError" in response["stderr"]


@pytest.mark.asyncio
async def test_bridge_server_not_running(
    execute_python_tool,
    mock_bridge_server
):
    """Test that error is raised if bridge server not running."""
    mock_bridge_server.is_running = Mock(return_value=False)
    
    arguments = {
        "code": "repos = github_search_repos(query='test')",
        "mcp_tools": ["github_search_repos"]
    }
    
    with pytest.raises(RuntimeError, match="bridge server"):
        await execute_python_tool.execute(arguments)