from uuid import uuid4

import pytest
from langgraph.checkpoint.memory import MemorySaver

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


@pytest.mark.asyncio
async def test_simple_agent():
    """Tests the simple agent."""
    agent = get_agent("simple")(
        name="Test Simple",
        instructions="Test instructions",
        model="anthropic/claude-haiku-4-5",
    )
    result = await agent.invoke(user_input="What is the capital of France?")
    assert result is not None
    last_message = result["messages"][-1]
    last_message_text = get_message_text(last_message)
    assert "paris" in last_message_text.lower()


@pytest.mark.asyncio
async def test_codeact_single_turn():
    """Tests the codeact-repl agent."""
    agent = get_agent("codeact-repl")(
        name="Test Codeact Repl",
        instructions="Test instructions",
        model="anthropic/claude-haiku-4-5",
        registry=AgentrRegistry(),
    )
    result = await agent.invoke(user_input="What is 2+2?")
    assert result is not None
    last_message = result["messages"][-1]
    last_message_text = get_message_text(last_message)
    assert "4" in last_message_text.lower()


@pytest.mark.asyncio
async def test_codeact_multi_turn():
    """Tests the codeact-repl agent."""
    checkpoint_saver = MemorySaver()
    async with get_agent("codeact-repl")(
        name="Test Codeact Repl",
        instructions="You are a helpful assistant",
        model="anthropic/claude-haiku-4-5",
        registry=AgentrRegistry(),
        memory=checkpoint_saver,
    ) as agent:
        thread_id = str(uuid4())
        result = await agent.invoke(
            user_input="Generate a function to calculate fibonnaci number, and get 10th number in the sequence. Use fib(0) = 0 and fib(1) = 1 as the base cases. Set x = fib(10)",
            thread_id=thread_id,
        )
        assert result is not None
        last_message = result["messages"][-1]
        last_message_text = get_message_text(last_message)
        assert "55" in last_message_text.lower()
        turn2 = await agent.invoke(
            user_input="What is the x+5?",
            thread_id=thread_id,
        )
        assert turn2 is not None
        last_message2 = turn2["messages"][-1]
        last_message2_text = get_message_text(last_message2)
        assert "60" in last_message2_text.lower()


@pytest.mark.asyncio
async def test_codeact_zenquotes_tool_workflow():
    """Tests the codeact agent with zenquotes tool workflow.

    Validates that:
    1. search_functions is called to find zenquotes tools
    2. load_functions is called to load the zenquotes tool
    3. execute_python_code is called to fetch the quote
    4. selected_tool_ids contains 'zenquotes__get_random_quote'
    """
    agent = get_agent("codeact-repl")(
        name="Test Codeact Zenquotes",
        instructions="Be very concise in your answers.",
        model="anthropic/claude-haiku-4-5",
        registry=AgentrRegistry(),
    )
    result = await agent.invoke(user_input="Find a quote with zenquotes")
    assert result is not None

    # Check that selected_tool_ids contains zenquotes__get_random_quote
    assert "selected_tool_ids" in result
    assert "zenquotes__get_random_quote" in result["selected_tool_ids"]

    # Check that the workflow includes the expected tool calls
    messages = result["messages"]
    tool_calls_made = set()

    for message in messages:
        if hasattr(message, "tool_calls") and message.tool_calls:
            for tool_call in message.tool_calls:
                tool_calls_made.add(tool_call["name"])

    # Validate that search_functions, load_functions, and execute_python_code were called
    assert "search_functions" in tool_calls_made, "search_functions should be called"
    assert "load_functions" in tool_calls_made, "load_functions should be called"
    assert "execute_python_code" in tool_calls_made, "execute_python_code should be called"

    # Verify that the last message contains a quote response
    last_message = result["messages"][-1]
    last_message_text = get_message_text(last_message)
    assert len(last_message_text) > 0, "Should return a quote response"
