"""
Tests for agent builder functionality with local file-based registry.

Tests the plan_agent and schedule_code_and_save_agent tools with LocalAgentBuilderRegistry.
"""

import shutil
from uuid import uuid4

import pytest
from langgraph.checkpoint.memory import MemorySaver

from universal_mcp.agentr.local_agent_builder_registry import LocalAgentBuilderRegistry
from universal_mcp.agentr.registry import AgentrRegistry
from universal_mcp.agents import get_agent
from universal_mcp.agents.utils import get_message_text


@pytest.fixture
def temp_agents_dir(tmp_path):
    """Create a temporary directory for agent storage."""
    agents_dir = tmp_path / "agents"
    agents_dir.mkdir()
    yield agents_dir
    # Cleanup
    if agents_dir.exists():
        shutil.rmtree(agents_dir)


@pytest.mark.asyncio
async def test_agent_builder_create_simple_agent(temp_agents_dir):
    """Test creating a simple agent using agent builder tools."""
    # Create local registry
    local_registry = LocalAgentBuilderRegistry(base_path=temp_agents_dir)
    checkpoint_saver = MemorySaver()
    thread_id = str(uuid4())

    # Create codeact agent with agent builder registry
    async with get_agent("codeact-repl")(
        name="Test Agent Builder",
        instructions="You are a helpful assistant that creates agents. Be very concise.",
        model="anthropic/claude-sonnet-4-5",
        registry=AgentrRegistry(),
        agent_builder_registry=local_registry,
        memory=checkpoint_saver,
    ) as agent:
        # Turn 1: Ask agent to create a simple calculator agent (will create plan and wait for approval)
        result1 = await agent.invoke(
            user_input="Create an agent that can add two numbers. Name it 'Calculator' and give it a simple description. The agent should accept two parameters a and b and return their sum.",
            thread_id=thread_id,
        )

        assert result1 is not None
        assert "plan" in result1
        assert result1["plan"] is not None

        # Verify plan_agent was called
        tool_calls_made_turn1 = set()
        for message in result1["messages"]:
            if hasattr(message, "tool_calls") and message.tool_calls:
                for tool_call in message.tool_calls:
                    tool_calls_made_turn1.add(tool_call["name"])
        assert "plan_agent" in tool_calls_made_turn1

        # Turn 2: Approve the plan and save the agent
        result2 = await agent.invoke(
            user_input="Looks good! Please save the agent now.",
            thread_id=thread_id,
        )

        assert result2 is not None

        # Check that schedule_code_and_save_agent was called
        tool_calls_made_turn2 = set()
        for message in result2["messages"]:
            if hasattr(message, "tool_calls") and message.tool_calls:
                for tool_call in message.tool_calls:
                    tool_calls_made_turn2.add(tool_call["name"])
        assert "schedule_code_and_save_agent" in tool_calls_made_turn2

    # Check that an agent was created in the local directory
    agent_dirs = list(temp_agents_dir.iterdir())
    assert len(agent_dirs) == 1, f"Expected 1 agent directory, found {len(agent_dirs)}"

    agent_dir = agent_dirs[0]
    assert (agent_dir / "metadata.json").exists()
    assert (agent_dir / "code.py").exists()

    # Read and verify the saved agent
    saved_agent = await local_registry.get_agent(agent_dir.name)
    assert saved_agent is not None
    assert saved_agent.name == "Calculator"
    assert "add" in saved_agent.description.lower() or "sum" in saved_agent.description.lower()
    assert saved_agent.instructions is not None
    assert "script" in saved_agent.instructions
    assert "def" in saved_agent.instructions["script"]  # Should contain Python function definition


@pytest.mark.asyncio
async def test_agent_builder_update_existing_agent(temp_agents_dir):
    """Test updating an existing agent using agent builder tools."""
    # Create an initial agent
    local_registry = LocalAgentBuilderRegistry(base_path=temp_agents_dir)
    initial_agent = await local_registry.create_agent(
        name="Simple Agent",
        description="A simple agent",
        instructions={
            "script": "def hello():\n    return 'Hello'",
            "plan": ["Say hello"],
        },
        tools={},
    )

    # Set the default agent ID so updates will target this agent
    local_registry.default_agent_id = initial_agent.id

    checkpoint_saver = MemorySaver()
    thread_id = str(uuid4())

    # Create codeact agent with agent builder registry
    async with get_agent("codeact-repl")(
        name="Test Agent Builder",
        instructions="You are a helpful assistant that creates agents. Be very concise.",
        model="anthropic/claude-sonnet-4-5",
        registry=AgentrRegistry(),
        agent_builder_registry=local_registry,
        memory=checkpoint_saver,
    ) as agent:
        # Turn 1: Ask agent to update the agent (will create/update plan and wait for approval)
        result1 = await agent.invoke(
            user_input="Update the agent to also have a goodbye function that returns 'Goodbye'.",
            thread_id=thread_id,
        )

        assert result1 is not None

        # Turn 2: Approve and save
        result2 = await agent.invoke(
            user_input="Looks good! Please save the updated agent.",
            thread_id=thread_id,
        )

        assert result2 is not None

    # Check that the agent was updated (not created)
    agent_dirs = list(temp_agents_dir.iterdir())
    assert len(agent_dirs) == 1, "Should still have only 1 agent (updated, not created new)"

    # Read and verify the updated agent
    updated_agent = await local_registry.get_agent(initial_agent.id)
    assert updated_agent is not None
    assert "goodbye" in updated_agent.instructions["script"].lower()
    assert "hello" in updated_agent.instructions["script"].lower()  # Should still have hello


@pytest.mark.asyncio
async def test_agent_builder_with_tools(temp_agents_dir):
    """Test creating an agent that uses external tools."""
    local_registry = LocalAgentBuilderRegistry(base_path=temp_agents_dir)
    checkpoint_saver = MemorySaver()
    thread_id = str(uuid4())

    async with get_agent("codeact-repl")(
        name="Test Agent Builder",
        instructions="You are a helpful assistant that creates agents. Be very concise.",
        model="anthropic/claude-sonnet-4-5",
        registry=AgentrRegistry(),
        agent_builder_registry=local_registry,
        memory=checkpoint_saver,
    ) as agent:
        # Turn 1: Ask agent to create a quote fetcher agent
        result1 = await agent.invoke(
            user_input="Create an agent called 'Quote Fetcher' that fetches random quotes using zenquotes. Make sure to search for and load the zenquotes tool.",
            thread_id=thread_id,
        )

        assert result1 is not None

        # Turn 2: Approve and save
        result2 = await agent.invoke(
            user_input="Looks good! Please save the agent.",
            thread_id=thread_id,
        )

        assert result2 is not None

    # Check that an agent was created
    agent_dirs = list(temp_agents_dir.iterdir())
    assert len(agent_dirs) == 1

    # Read and verify the saved agent
    saved_agent = await local_registry.get_agent(agent_dirs[0].name)
    assert saved_agent is not None
    assert saved_agent.name == "Quote Fetcher"
    assert saved_agent.tools is not None
    assert "zenquotes" in str(saved_agent.tools).lower() or "quote" in saved_agent.instructions["script"].lower()


@pytest.mark.asyncio
async def test_agent_builder_plan_modification(temp_agents_dir):
    """Test that agent builder can modify plans across turns."""
    local_registry = LocalAgentBuilderRegistry(base_path=temp_agents_dir)
    checkpoint_saver = MemorySaver()

    async with get_agent("codeact-repl")(
        name="Test Agent Builder",
        instructions="You are a helpful assistant that creates agents. Be very concise.",
        model="anthropic/claude-sonnet-4-5",
        registry=AgentrRegistry(),
        agent_builder_registry=local_registry,
        memory=checkpoint_saver,
    ) as agent:
        thread_id = str(uuid4())

        # First turn: Create initial plan for a simple math agent
        result1 = await agent.invoke(
            user_input="Create a plan for an agent that can calculate fibonacci numbers. Just create the plan, don't save yet.",
            thread_id=thread_id,
        )

        assert result1 is not None
        # Should have plan in state
        assert "plan" in result1
        assert result1["plan"] is not None
        assert len(result1["plan"]) > 0

        # Second turn: Modify the plan
        result2 = await agent.invoke(
            user_input="Update the plan to also calculate factorial in addition to fibonacci.",
            thread_id=thread_id,
        )

        assert result2 is not None
        assert "plan" in result2
        # Plan should now mention factorial
        plan_text = "\n".join(result2["plan"])
        assert "factorial" in plan_text.lower()

        # Third turn: Save the agent
        result3 = await agent.invoke(
            user_input="Now save the agent with name 'Math Helper'.",
            thread_id=thread_id,
        )

        assert result3 is not None

    # Verify the agent was saved with the updated plan
    agent_dirs = list(temp_agents_dir.iterdir())
    assert len(agent_dirs) == 1

    saved_agent = await local_registry.get_agent(agent_dirs[0].name)
    assert saved_agent is not None
    assert saved_agent.name == "Math Helper"
    assert "plan" in saved_agent.instructions
    plan_text = "\n".join(saved_agent.instructions["plan"])
    assert "factorial" in plan_text.lower()


@pytest.mark.asyncio
async def test_agent_builder_create_and_load_by_id(temp_agents_dir):
    """
    Test creating an agent with agent builder, saving it, then loading it by agent_id
    and verifying that the functions defined in the agent work correctly.
    """
    # Step 1: Create and save an agent using agent builder
    local_registry = LocalAgentBuilderRegistry(base_path=temp_agents_dir)
    checkpoint_saver = MemorySaver()
    thread_id = str(uuid4())

    # Create codeact agent with agent builder registry
    async with get_agent("codeact-repl")(
        name="Test Agent Builder",
        instructions="You are a helpful assistant that creates agents. Be very concise.",
        model="anthropic/claude-sonnet-4-5",
        registry=AgentrRegistry(),
        agent_builder_registry=local_registry,
        memory=checkpoint_saver,
    ) as builder_agent:
        # Turn 1: Ask agent to create a simple temperature converter agent
        result1 = await builder_agent.invoke(
            user_input="Create an agent called 'Temperature Converter' that has a function to convert Celsius to Fahrenheit. The function should be named 'celsius_to_fahrenheit' and take a parameter 'celsius' and return the fahrenheit value using the formula: fahrenheit = (celsius * 9/5) + 32",
            thread_id=thread_id,
        )

        assert result1 is not None
        assert "plan" in result1
        assert result1["plan"] is not None

        # Turn 2: Approve and save the agent
        result2 = await builder_agent.invoke(
            user_input="Looks good! Please save the agent now.",
            thread_id=thread_id,
        )

        assert result2 is not None

    # Verify agent was created
    agent_dirs = list(temp_agents_dir.iterdir())
    assert len(agent_dirs) == 1, f"Expected 1 agent directory, found {len(agent_dirs)}"

    agent_id = agent_dirs[0].name

    # Verify the saved agent has the correct structure
    saved_agent_metadata = await local_registry.get_agent(agent_id)
    assert saved_agent_metadata is not None
    assert saved_agent_metadata.name == "Temperature Converter"
    assert "script" in saved_agent_metadata.instructions
    script_content = saved_agent_metadata.instructions["script"]
    assert "celsius_to_fahrenheit" in script_content
    assert "def" in script_content

    # Step 2: Create a NEW agent instance that loads the saved agent by agent_id
    new_registry = LocalAgentBuilderRegistry(base_path=temp_agents_dir, default_agent_id=agent_id)
    new_checkpoint = MemorySaver()
    new_thread_id = str(uuid4())

    async with get_agent("codeact-repl")(
        name="Test Loaded Agent",
        instructions="You are a helpful assistant. Be very concise.",
        model="anthropic/claude-haiku-4-5",
        registry=AgentrRegistry(),
        agent_builder_registry=new_registry,
        memory=new_checkpoint,
    ) as loaded_agent:
        # Step 3: Test that the loaded agent can use the celsius_to_fahrenheit function
        result3 = await loaded_agent.invoke(
            user_input="Convert 0 degrees Celsius to Fahrenheit using the celsius_to_fahrenheit function",
            thread_id=new_thread_id,
        )

        assert result3 is not None
        messages = result3["messages"]
        assert len(messages) > 0

        # Check that the function was executed
        last_message_text = get_message_text(messages[-1])
        # 0°C = 32°F
        assert "32" in last_message_text, f"Expected '32' in response, got: {last_message_text}"

        # Step 4: Test with a different value to ensure the function works correctly
        result4 = await loaded_agent.invoke(
            user_input="Now convert 100 degrees Celsius to Fahrenheit",
            thread_id=new_thread_id,
        )

        assert result4 is not None
        last_message_text_2 = get_message_text(result4["messages"][-1])
        # 100°C = 212°F
        assert "212" in last_message_text_2, f"Expected '212' in response, got: {last_message_text_2}"
