"""Tests for SubAgent support (AgentTool)."""
import pytest
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4

from yier_agents.src.agent import Agent, AgentResult, AgentTool
from yier_agents.src.message import Message
from yier_agents.src.tool import Tool, ToolContext, ToolOutput
from yier_agents.src.llm import LLM
from pydantic import BaseModel


class SimpleParams(BaseModel):
    """Simple tool parameters for testing."""
    query: str


async def simple_execute(params: SimpleParams, ctx: ToolContext) -> ToolOutput:
    """Simple tool executor for testing."""
    return ToolOutput(content=f"Result: {params.query}")


@pytest.mark.asyncio
async def test_agent_tool_creation():
    """Test AgentTool creation."""
    # Create mock LLM and sub-agent
    mock_llm = MagicMock(spec=LLM)
    sub_agent = Agent(llm=mock_llm, tools=[], system_prompt="Sub-agent system")

    # Create AgentTool
    agent_tool = AgentTool(
        name="research_tool",
        description="Research tool using sub-agent",
        agent=sub_agent,
        inherit_context=True
    )

    # Verify basic properties
    assert agent_tool.name == "research_tool"
    assert agent_tool.description == "Research tool using sub-agent"
    assert agent_tool.agent == sub_agent

    # Verify OpenAI schema
    schema = agent_tool.to_openai_schema()
    assert schema["type"] == "function"
    assert schema["function"]["name"] == "research_tool"
    assert "prompt" in schema["function"]["parameters"]["properties"]


@pytest.mark.asyncio
async def test_agent_tool_execute():
    """Test AgentTool execution."""
    # Create mock LLM
    mock_llm = AsyncMock(spec=LLM)

    # Create mock sub-agent response
    from yier_agents.src.message import ToolCall
    mock_result = AgentResult(
        session_id="sub-session-123",
        messages=[
            Message(role="user", content="Research question"),
            Message(role="assistant", content="Research answer", tool_calls=[
                ToolCall(id="tool-1", name="search", arguments={"query": "test"})
            ]),
            Message(role="tool", content="Search result", tool_call_id="tool-1"),
            Message(role="assistant", content="Final answer based on search")
        ],
        final_message=Message(role="assistant", content="Final answer based on search"),
        finish_reason="stop"
    )

    # Create sub-agent with mocked run method
    sub_agent = Agent(llm=mock_llm, tools=[], system_prompt="Sub-agent")
    sub_agent.run = AsyncMock(return_value=mock_result)

    # Create AgentTool
    agent_tool = AgentTool(
        name="research_tool",
        description="Research tool",
        agent=sub_agent,
        inherit_context=True
    )

    # Execute the tool
    from yier_agents.src.agent import AgentToolParams
    params = AgentToolParams(prompt="Research something")
    ctx = ToolContext(session_id="main-123", message_id="msg-456", call_id="call-789")

    result = await agent_tool.execute(params, ctx)

    # Verify result
    assert isinstance(result, ToolOutput)
    assert "Final answer based on search" in result.content
    assert "<subagent_metadata>" in result.content
    assert "search(1)" in result.content  # Tool summary
    assert "sub-session-123" in result.content

    # Verify sub-agent was called with correct session ID
    sub_agent.run.assert_called_once()
    call_args = sub_agent.run.call_args
    assert call_args[0][0] == "Research something"  # prompt
    assert call_args[0][1].startswith("main-123:sub:")  # session_id format


@pytest.mark.asyncio
async def test_main_agent_with_subagent():
    """Test main agent calling sub-agent through AgentTool."""
    # Create mock LLM for main agent
    main_llm = AsyncMock(spec=LLM)

    # Create mock sub-agent
    sub_llm = AsyncMock(spec=LLM)
    sub_agent = Agent(llm=sub_llm, tools=[], system_prompt="Sub-agent")

    # Mock sub-agent response
    sub_result = AgentResult(
        session_id="sub-session-456",
        messages=[
            Message(role="user", content="Sub task"),
            Message(role="assistant", content="Sub result")
        ],
        final_message=Message(role="assistant", content="Sub result"),
        finish_reason="stop"
    )
    sub_agent.run = AsyncMock(return_value=sub_result)

    # Create AgentTool
    agent_tool = AgentTool(
        name="delegate_tool",
        description="Delegate to sub-agent",
        agent=sub_agent,
        inherit_context=True
    )

    # Create main agent with AgentTool
    main_agent = Agent(llm=main_llm, tools=[agent_tool], system_prompt="Main agent")

    # Mock main agent LLM responses
    from yier_agents.src.message import LLMResponse, ToolCall

    # First response: call the sub-agent tool
    main_llm.complete.return_value = LLMResponse(
        message=Message(
            role="assistant",
            content=None,
            tool_calls=[
                ToolCall(id="tc-1", name="delegate_tool", arguments={"prompt": "Do research"})
            ]
        ),
        finish_reason="tool_calls"
    )

    # Execute main agent (need to manually handle the flow since we're mocking)
    # This is a simplified integration test
    from yier_agents.src.agent import AgentToolParams
    params = AgentToolParams(prompt="Do research")
    ctx = ToolContext(session_id="main-session", message_id="msg-1", call_id="tc-1")

    result = await agent_tool.execute(params, ctx)

    # Verify the tool was executed and returned expected content
    assert isinstance(result, ToolOutput)
    assert "Sub result" in result.content

    # Verify sub-agent was called
    sub_agent.run.assert_called_once()
    assert sub_agent.run.call_args[0][0] == "Do research"
    assert sub_agent.run.call_args[0][1].startswith("main-session:sub:")
