"""
Parametrized tests for sandbox implementations.

Tests all sandbox implementations (InProcessSandbox, SubprocessSandbox, E2BSandbox)
with the same test cases to ensure consistent behavior across different execution environments.
"""

import pytest

from universal_mcp.agents.Sandbox.e2b_sandbox import E2BSandbox
from universal_mcp.agents.Sandbox.in_process_sandbox import InProcessSandbox
from universal_mcp.agents.Sandbox.subprocess_sandbox import SubprocessSandbox


# Parametrize tests across different sandbox implementations
@pytest.fixture(
    params=[
        pytest.param("in_process", id="InProcessSandbox"),
        pytest.param("subprocess", id="SubprocessSandbox"),
        pytest.param("e2b", id="E2BSandbox", marks=pytest.mark.skip(reason="E2B requires API key")),
    ]
)
def sandbox(request):
    """Create a sandbox instance based on the parameter."""
    sandbox_type = request.param
    timeout = 10

    if sandbox_type == "in_process":
        sbx = InProcessSandbox(timeout=timeout)
    elif sandbox_type == "subprocess":
        sbx = SubprocessSandbox(timeout=timeout)
    elif sandbox_type == "e2b":
        sbx = E2BSandbox(timeout=timeout)
    else:
        raise ValueError(f"Unknown sandbox type: {sandbox_type}")

    yield sbx

    # Cleanup
    if hasattr(sbx, "__del__"):
        try:
            sbx.__del__()
        except Exception:
            pass


@pytest.mark.asyncio
async def test_simple_exec(sandbox):
    """Test basic code execution."""
    result = await sandbox.run("print('Hello World')")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "Hello World", f"Expected 'Hello World', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error_type, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_multi_turn(sandbox):
    """Test context persistence across multiple executions."""
    await sandbox.run("a = 5")
    await sandbox.run("b = 6")
    result = await sandbox.run("print(a + b)")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "11", f"Expected '11', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_update_context_with_variable(sandbox):
    """Test injecting variables into sandbox context."""
    await sandbox.update_context({"a": 5})
    result = await sandbox.run("print(a)")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "5", f"Expected '5', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_update_context_with_function(sandbox):
    """Test injecting functions into sandbox context."""

    def hello():
        print("Hello World")  # noqa: T201

    await sandbox.update_context({"hello": hello})
    result = await sandbox.run("hello()")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "Hello World", f"Expected 'Hello World', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_update_context_with_class(sandbox):
    """Test injecting classes into sandbox context."""

    class Hello:
        def __init__(self):
            pass

        def world(self):
            return "World"

    await sandbox.update_context({"Hello": Hello})
    result = await sandbox.run("print(Hello().world())")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "World", f"Expected 'World', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "sandbox_class",
    [
        pytest.param(InProcessSandbox, id="InProcessSandbox"),
        pytest.param(SubprocessSandbox, id="SubprocessSandbox"),
    ],
)
async def test_cross_persistence(sandbox_class):
    """Test context serialization and restoration across sandbox instances."""
    sandbox1 = sandbox_class(timeout=10)
    await sandbox1.run("a = 5")
    context = await sandbox1.get_context()

    # Create new sandbox and restore context
    sandbox2 = sandbox_class(timeout=10)
    await sandbox2.update_context(context)
    result = await sandbox2.run("print(a)")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "5", f"Expected '5', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_syntax_error(sandbox):
    """Test that syntax errors are handled gracefully."""
    result = await sandbox.run("print('Hello World'")  # Missing closing parenthesis
    assert result["exit_code"] != 0, "Expected non-zero exit code for syntax error"
    assert result["error_type"] in ["SyntaxError", "IndentationError"], (
        f"Expected SyntaxError, got {result['error_type']}"
    )
    assert result["error_message"] is not None and "Recovery" in result["error_message"], "Expected recovery hints"


@pytest.mark.asyncio
async def test_name_error(sandbox):
    """Test that undefined variables raise NameError."""
    result = await sandbox.run("print(undefined_variable)")
    assert result["exit_code"] != 0, "Expected non-zero exit code for NameError"
    assert result["error_type"] == "NameError", f"Expected NameError, got {result['error_type']}"
    assert result["error_message"] is not None and "not defined" in result["error_message"], (
        "Expected 'not defined' in error message"
    )


@pytest.mark.asyncio
async def test_import_error(sandbox):
    """Test that missing imports are handled."""
    result = await sandbox.run("import nonexistent_module")
    assert result["exit_code"] != 0, "Expected non-zero exit code for ImportError"
    assert result["error_type"] in ["ImportError", "ModuleNotFoundError"], (
        f"Expected ImportError or ModuleNotFoundError, got {result['error_type']}"
    )
    assert result["error_message"] is not None and "Recovery" in result["error_message"], "Expected recovery hints"


@pytest.mark.asyncio
async def test_zero_division_error(sandbox):
    """Test that division by zero is handled."""
    result = await sandbox.run("result = 1 / 0")
    assert result["exit_code"] != 0, "Expected non-zero exit code for ZeroDivisionError"
    assert result["error_type"] == "ZeroDivisionError", f"Expected ZeroDivisionError, got {result['error_type']}"
    assert result["error_message"] is not None and "Recovery" in result["error_message"], "Expected recovery hints"


@pytest.mark.asyncio
async def test_complex_context_preservation(sandbox):
    """Test preservation of complex data structures."""
    await sandbox.run("data = {'key': 'value', 'numbers': [1, 2, 3]}")
    result = await sandbox.run("print(data['key'])")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "value", f"Expected 'value', got '{result['stdout']}'"

    result = await sandbox.run("print(sum(data['numbers']))")
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "6", f"Expected '6', got '{result['stdout']}'"


@pytest.mark.asyncio
async def test_multiple_prints(sandbox):
    """Test capturing multiple print statements."""
    code = """
print('Line 1')
print('Line 2')
print('Line 3')
"""
    result = await sandbox.run(code)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    lines = result["stdout"].strip().split("\n")
    assert len(lines) == 3, f"Expected 3 lines, got {len(lines)}"
    assert lines[0] == "Line 1"
    assert lines[1] == "Line 2"
    assert lines[2] == "Line 3"


@pytest.mark.asyncio
async def test_async_code_execution(sandbox):
    """Test execution of async code."""
    code = """
async def async_func():
    return "async result"

result = await async_func()
print(result)
"""
    result = await sandbox.run(code)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "async result", f"Expected 'async result', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_user_defined_function_persistence(sandbox):
    """Test that user-defined functions persist across executions."""
    # Define a function in first execution
    code1 = """
def calculate_square(x):
    return x * x

result = calculate_square(5)
print(result)
"""
    result = await sandbox.run(code1)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "25", f"Expected '25', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"

    # Use the same function in second execution (should persist)
    code2 = """
result = calculate_square(7)
print(result)
"""
    result = await sandbox.run(code2)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "49", f"Expected '49', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_user_defined_class_persistence(sandbox):
    """Test that user-defined classes persist across executions."""
    # Define a class in first execution
    code1 = """
class Calculator:
    def __init__(self, value):
        self.value = value

    def add(self, x):
        return self.value + x

calc = Calculator(10)
print(calc.add(5))
"""
    result = await sandbox.run(code1)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "15", f"Expected '15', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"

    # Use the same class in second execution (should persist)
    code2 = """
calc2 = Calculator(20)
print(calc2.add(15))
"""
    result = await sandbox.run(code2)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "35", f"Expected '35', got '{result['stdout']}'"
    assert result["error_type"] is None, f"Expected no error, got '{result['error_type']}'"


@pytest.mark.asyncio
async def test_builtin_exclusion(sandbox):
    """Test that built-in functions are not polluted by overriding."""
    # Override a builtin in first execution
    code1 = """
# Try to override a builtin (should not persist)
len_original = len
len = lambda x: 999
print(len([1, 2, 3]))
"""
    result = await sandbox.run(code1)
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"
    assert result["stdout"].strip() == "999", f"Expected '999', got '{result['stdout']}'"

    # In second execution, len should still be the builtin
    code2 = """
print(len([1, 2, 3, 4]))
"""
    result = await sandbox.run(code2)
    # This might be 999 if our override persisted, or 4 if builtin was restored
    # The expected behavior is implementation-specific
    # For now, just verify it doesn't crash
    assert result["exit_code"] == 0, f"Expected exit_code 0, got {result['exit_code']}"


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