"""
Test context persistence across turns with cloudpickle serialization.

This tests that functions, classes, and imports persist correctly
when context is serialized as base64-encoded cloudpickle strings.
"""

import pytest
from langgraph.checkpoint.memory import MemorySaver

from universal_mcp.agentr.registry import AgentrRegistry
from universal_mcp.agents.codeact0.agent import CodeActPlaybookAgent


@pytest.mark.asyncio
async def test_variable_persistence():
    """Test that variables persist across turns."""
    memory = MemorySaver()
    async with CodeActPlaybookAgent(
        name="Variable Test Agent",
        instructions="Execute code exactly as provided.",
        model="anthropic:claude-haiku-4-5",
        registry=AgentrRegistry(),
        memory=memory,
    ) as agent:
        # Turn 1: Define variables
        _ = await agent.invoke(
            user_input="""
x = 42
my_list = [1, 2, 3]
my_dict = {"name": "test", "value": 100}
smart_print(f"Defined: x={x}, my_list={my_list}, my_dict={my_dict}")
            """,
            thread_id="test_vars",
        )

        # Turn 2: Use variables from Turn 1
        result2 = await agent.invoke(
            user_input="""
# Verify variables are available
smart_print(f"x = {x}")
smart_print(f"my_list = {my_list}")
smart_print(f"my_dict = {my_dict}")

# Modify them
x = x + 10
my_list.append(4)
smart_print(f"Modified: x={x}, my_list={my_list}")
            """,
            thread_id="test_vars",
        )

        # Check that no errors occurred
        last_msg = str(result2["messages"][-1].content)
        assert "Error" not in last_msg, f"Variables not persisted: {last_msg}"
        assert "NameError" not in last_msg, f"Variables not persisted: {last_msg}"


@pytest.mark.asyncio
async def test_function_persistence():
    """Test that user-defined functions persist across turns."""
    memory = MemorySaver()
    async with CodeActPlaybookAgent(
        name="Function Test Agent",
        instructions="Execute code exactly as provided.",
        model="anthropic:claude-haiku-4-5",
        registry=AgentrRegistry(),
        memory=memory,
    ) as agent:
        # Turn 1: Define an async function
        _ = await agent.invoke(
            user_input="""
async def double(x):
    return x * 2

result = await double(21)
smart_print(f"double(21) = {result}")
            """,
            thread_id="test_func",
        )

        # Turn 2: Use function from Turn 1
        result2 = await agent.invoke(
            user_input="""
# Use function defined in Turn 1
result2 = await double(50)
smart_print(f"double(50) = {result2}")
            """,
            thread_id="test_func",
        )

        # Check that function was available
        last_msg = str(result2["messages"][-1].content)
        assert "NameError" not in last_msg, f"Function not persisted: {last_msg}"
        assert "not defined" not in last_msg, f"Function not persisted: {last_msg}"


@pytest.mark.asyncio
async def test_class_persistence():
    """Test that user-defined classes persist across turns."""
    memory = MemorySaver()
    async with CodeActPlaybookAgent(
        name="Class Test Agent",
        instructions="Execute code exactly as provided.",
        model="anthropic:claude-haiku-4-5",
        registry=AgentrRegistry(),
        memory=memory,
    ) as agent:
        # Turn 1: Define a class
        _ = await agent.invoke(
            user_input="""
class Calculator:
    def __init__(self, value):
        self.value = value

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

# Note: We only save the class definition, not instances
smart_print("Calculator class defined")
            """,
            thread_id="test_class",
        )

        # Turn 2: Use class from Turn 1
        result2 = await agent.invoke(
            user_input="""
# Create new instance of class from Turn 1
calc = Calculator(20)
result = calc.add(5)
smart_print(f"calc.value = {calc.value}")
            """,
            thread_id="test_class",
        )

        # Check that class was available
        last_msg = str(result2["messages"][-1].content)
        assert "NameError" not in last_msg, f"Class not persisted: {last_msg}"
        assert "not defined" not in last_msg, f"Class not persisted: {last_msg}"


@pytest.mark.asyncio
async def test_import_persistence():
    """Test that imports persist across turns."""
    memory = MemorySaver()
    async with CodeActPlaybookAgent(
        name="Import Test Agent",
        instructions="Execute code exactly as provided.",
        model="anthropic:claude-haiku-4-5",
        registry=AgentrRegistry(),
        memory=memory,
    ) as agent:
        # Turn 1: Import module
        _ = await agent.invoke(
            user_input="""
import math

result = math.sqrt(16)
smart_print(f"math.sqrt(16) = {result}")
smart_print(f"math.pi = {math.pi}")
            """,
            thread_id="test_import",
        )

        # Turn 2: Use import from Turn 1
        result2 = await agent.invoke(
            user_input="""
# Use math module from Turn 1
result = math.sqrt(25)
smart_print(f"math.sqrt(25) = {result}")
            """,
            thread_id="test_import",
        )

        # Check that import was available
        last_msg = str(result2["messages"][-1].content)
        assert "NameError" not in last_msg, f"Import not persisted: {last_msg}"
        assert "not defined" not in last_msg, f"Import not persisted: {last_msg}"


@pytest.mark.asyncio
async def test_combined_persistence():
    """Test that functions, classes, and imports persist across agent instances with same memory."""
    # Create shared memory
    memory = MemorySaver()
    registry = AgentrRegistry()

    # Agent 1: Define function with import
    async with CodeActPlaybookAgent(
        name="Combined Test Agent 1",
        instructions="Execute code exactly as provided.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
    ) as agent1:
        # Turn 1: Define function with import
        _ = await agent1.invoke(
            user_input="""
import math

async def calculate_circle_area(radius):
    return math.pi * radius ** 2

area = await calculate_circle_area(5)
smart_print(f"Area of circle with radius 5: {area}")
            """,
            thread_id="test_combined",
        )

    # Agent 2: Create NEW agent instance with SAME memory and thread_id
    async with CodeActPlaybookAgent(
        name="Combined Test Agent 2",
        instructions="Execute code exactly as provided.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
    ) as agent2:
        # Turn 2: Use both function and import from Turn 1 (different agent instance!)
        result2 = await agent2.invoke(
            user_input="""
# Use function and import from Turn 1 (created in different agent instance)
area2 = await calculate_circle_area(10)
smart_print(f"Area of circle with radius 10: {area2}")

# Also verify we can use math directly
circumference = 2 * math.pi * 10
smart_print(f"Circumference: {circumference}")
            """,
            thread_id="test_combined",
        )

        # Check that both function and import were available
        last_msg = str(result2["messages"][-1].content)
        assert "NameError" not in last_msg, f"Function/import not persisted across agent instances: {last_msg}"
        assert "not defined" not in last_msg, f"Function/import not persisted across agent instances: {last_msg}"
