"""
Test that registry tools loaded via load_functions persist across turns.

This tests that tools loaded from the registry (like google_mail__list_messages)
are available across multiple conversation turns when serialized with cloudpickle.
"""

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_registry_tool_persistence():
    """Test that tools loaded via load_functions persist across turns."""
    memory = MemorySaver()
    registry = AgentrRegistry()
    async with CodeActPlaybookAgent(
        name="Registry Tool Test Agent",
        instructions="Execute code and use loaded tools as requested.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
    ) as agent:
        # Turn 1: Search for and load a tool
        _ = await agent.invoke(
            user_input="""
Search for google mail tools and load the list_messages function.
Then verify it's available by checking if it's a callable function.
            """,
            thread_id="test_registry_tools",
        )

        # Turn 2: Verify the tool is still available without reloading
        result2 = await agent.invoke(
            user_input="""
Check if google_mail__list_messages is still available in the namespace.
Print its type and confirm it's callable.
DO NOT reload the function - just check if it exists from the previous turn.
            """,
            thread_id="test_registry_tools",
        )

        # Check that the tool is still available
        last_msg = str(result2["messages"][-1].content)

        # The test should NOT have NameError (tool should be persisted)
        assert "NameError" not in last_msg, f"Tool not persisted: {last_msg}"
        assert "not defined" not in last_msg, f"Tool not persisted: {last_msg}"


@pytest.mark.asyncio
async def test_loaded_tool_execution_across_turns():
    """Test that loaded tools can be executed across turns."""
    memory = MemorySaver()
    registry = AgentrRegistry()
    async with CodeActPlaybookAgent(
        name="Tool Execution Test Agent",
        instructions="Execute code and use loaded tools as requested.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
    ) as agent:
        # Turn 1: Load LLM tools
        _ = await agent.invoke(
            user_input="""
Load the llm__generate_text function.
Store a simple test string in a variable called 'test_input' with value "Hello world".
            """,
            thread_id="test_tool_execution",
        )

        # Turn 2: Use the loaded tool from Turn 1
        result2 = await agent.invoke(
            user_input="""
Use the llm__generate_text function that was loaded in the previous turn.
Use it to generate a greeting message using the 'test_input' variable.
            """,
            thread_id="test_tool_execution",
        )

        # Check that the tool execution succeeded
        last_msg = str(result2["messages"][-1].content)

        # Should not have errors about missing function or variable
        assert "NameError" not in last_msg, f"Tool or variable not persisted: {last_msg}"
        assert "not defined" not in last_msg, f"Tool or variable not persisted: {last_msg}"


@pytest.mark.asyncio
async def test_multiple_tools_persistence():
    """Test that multiple loaded tools persist together."""
    memory = MemorySaver()
    registry = AgentrRegistry()
    async with CodeActPlaybookAgent(
        name="Multiple Tools Test Agent",
        instructions="Execute code and use loaded tools as requested.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
    ) as agent:
        # Turn 1: Load multiple LLM tools
        _ = await agent.invoke(
            user_input="""
Load these LLM functions:
1. llm__generate_text
2. llm__classify_data
3. llm__extract_data

Store their names in a list called 'loaded_tools'.
            """,
            thread_id="test_multiple_tools",
        )

        # Turn 2: Verify all tools are still available
        result2 = await agent.invoke(
            user_input="""
Check if all three functions from Turn 1 are still available:
- llm__generate_text
- llm__classify_data
- llm__extract_data

Also verify the 'loaded_tools' list is still available.
Print the status of each.
            """,
            thread_id="test_multiple_tools",
        )

        # Check that all tools are still available
        last_msg = str(result2["messages"][-1].content)

        assert "NameError" not in last_msg, f"Tools not persisted: {last_msg}"
        assert "not defined" not in last_msg, f"Tools not persisted: {last_msg}"
