"""Integration tests for Agent system."""

import pytest
from unittest.mock import AsyncMock, MagicMock
from pydantic import BaseModel, Field

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


# Define calculator tool for integration tests
class CalculatorParams(BaseModel):
    """Parameters for calculator tool."""
    operation: str = Field(..., description="Operation: add, subtract, multiply, divide")
    x: float = Field(..., description="First number")
    y: float = Field(..., description="Second number")


async def calculator_execute(params: CalculatorParams, ctx: ToolContext) -> ToolOutput:
    """Execute calculator operation."""
    operations = {
        "add": lambda x, y: x + y,
        "subtract": lambda x, y: x - y,
        "multiply": lambda x, y: x * y,
        "divide": lambda x, y: x / y if y != 0 else None,
    }

    if params.operation not in operations:
        return ToolOutput(
            content=f"Error: Unknown operation '{params.operation}'",
            metadata={"error": True}
        )

    result = operations[params.operation](params.x, params.y)

    if result is None:
        return ToolOutput(
            content="Error: Division by zero",
            metadata={"error": True}
        )

    return ToolOutput(
        content=f"Result: {result}",
        metadata={"result": result, "operation": params.operation}
    )


@pytest.fixture
def calculator_tool():
    """Create a calculator tool."""
    return Tool(
        name="calculator",
        description="Perform basic arithmetic operations",
        parameters=CalculatorParams,
        execute=calculator_execute,
    )


@pytest.fixture
def mock_llm():
    """Create a mock LLM."""
    llm = MagicMock(spec=LLM)
    llm.complete = AsyncMock()
    return llm


@pytest.mark.asyncio
async def test_full_agent_workflow(mock_llm, calculator_tool):
    """Test complete agent workflow with calculator tool.

    This test simulates a real conversation:
    1. User asks to calculate 15 + 27
    2. Agent calls calculator tool
    3. Agent responds with result
    """
    # Setup mock LLM responses
    mock_llm.complete.side_effect = [
        # First call: LLM decides to use calculator
        LLMResponse(
            message=Message(
                role="assistant",
                content=None,
                tool_calls=[
                    ToolCall(
                        id="call-calc-1",
                        name="calculator",
                        arguments={"operation": "add", "x": 15, "y": 27}
                    )
                ]
            ),
            finish_reason="tool_calls"
        ),
        # Second call: LLM responds with final answer
        LLMResponse(
            message=Message(
                role="assistant",
                content="The result of 15 + 27 is 42."
            ),
            finish_reason="stop"
        )
    ]

    # Create agent
    agent = Agent(
        llm=mock_llm,
        tools=[calculator_tool],
        max_iterations=10,
        system_prompt="You are a helpful math assistant."
    )

    # Run agent
    result = await agent.run(
        user_message="What is 15 plus 27?",
        session_id="integration-test-1"
    )

    # Verify result
    assert isinstance(result, AgentResult)
    assert result.session_id == "integration-test-1"
    assert result.finish_reason == "stop"
    assert result.final_message.role == "assistant"
    assert "42" in result.final_message.content

    # Verify message flow
    # system + user + assistant(tool_call) + tool + assistant(final) = 5
    assert len(result.messages) == 5
    assert result.messages[0].role == "system"
    assert result.messages[1].role == "user"
    assert result.messages[2].role == "assistant"
    assert result.messages[2].tool_calls is not None
    assert result.messages[3].role == "tool"
    assert "Result: 42.0" in result.messages[3].content
    assert result.messages[4].role == "assistant"

    # Verify LLM was called twice
    assert mock_llm.complete.call_count == 2

    # Verify tool schemas were passed
    first_call_args = mock_llm.complete.call_args_list[0]
    assert "tools" in first_call_args[1]
    tools_schema = first_call_args[1]["tools"]
    assert len(tools_schema) == 1
    assert tools_schema[0]["function"]["name"] == "calculator"


@pytest.mark.asyncio
async def test_agent_with_subagent_integration(mock_llm, calculator_tool):
    """Test main agent calling sub-agent via AgentTool.

    This test simulates hierarchical agent delegation:
    1. Main agent receives user request
    2. Main agent delegates to math sub-agent via AgentTool
    3. Sub-agent uses calculator tool
    4. Sub-agent returns result to main agent
    5. Main agent responds to user
    """
    # Create sub-agent (math specialist)
    sub_llm = MagicMock(spec=LLM)
    sub_llm.complete = AsyncMock()

    # Sub-agent LLM responses
    sub_llm.complete.side_effect = [
        # Sub-agent calls calculator
        LLMResponse(
            message=Message(
                role="assistant",
                content=None,
                tool_calls=[
                    ToolCall(
                        id="sub-call-1",
                        name="calculator",
                        arguments={"operation": "multiply", "x": 6, "y": 7}
                    )
                ]
            ),
            finish_reason="tool_calls"
        ),
        # Sub-agent returns result
        LLMResponse(
            message=Message(
                role="assistant",
                content="The result of 6 × 7 is 42."
            ),
            finish_reason="stop"
        )
    ]

    sub_agent = Agent(
        llm=sub_llm,
        tools=[calculator_tool],
        max_iterations=10,
        system_prompt="You are a math calculation specialist."
    )

    # Create AgentTool wrapping sub-agent
    math_agent_tool = AgentTool(
        name="math_agent",
        description="Delegate math calculations to specialized sub-agent",
        agent=sub_agent,
        inherit_context=True
    )

    # Main agent LLM responses
    mock_llm.complete.side_effect = [
        # Main agent delegates to math sub-agent
        LLMResponse(
            message=Message(
                role="assistant",
                content=None,
                tool_calls=[
                    ToolCall(
                        id="main-call-1",
                        name="math_agent",
                        arguments={
                            "prompt": "Calculate 6 times 7",
                            "session_id": None
                        }
                    )
                ]
            ),
            finish_reason="tool_calls"
        ),
        # Main agent responds with final answer
        LLMResponse(
            message=Message(
                role="assistant",
                content="I delegated the calculation to my math specialist, and the answer is 42."
            ),
            finish_reason="stop"
        )
    ]

    # Create main agent with math_agent_tool
    main_agent = Agent(
        llm=mock_llm,
        tools=[math_agent_tool],
        max_iterations=10,
        system_prompt="You are a helpful assistant that can delegate tasks."
    )

    # Run main agent
    result = await main_agent.run(
        user_message="Can you help me calculate 6 times 7?",
        session_id="integration-test-2"
    )

    # Verify result
    assert isinstance(result, AgentResult)
    assert result.session_id == "integration-test-2"
    assert result.finish_reason == "stop"
    assert "42" in result.final_message.content

    # Verify message flow in main agent
    # system + user + assistant(tool_call) + tool(subagent_result) + assistant(final) = 5
    assert len(result.messages) == 5
    assert result.messages[0].role == "system"
    assert result.messages[1].role == "user"
    assert result.messages[2].role == "assistant"
    assert result.messages[2].tool_calls is not None
    assert result.messages[2].tool_calls[0].name == "math_agent"
    assert result.messages[3].role == "tool"
    assert "42" in result.messages[3].content
    assert "<subagent_metadata>" in result.messages[3].content
    assert "calculator(1)" in result.messages[3].content  # Tool usage summary
    assert result.messages[4].role == "assistant"

    # Verify both main and sub-agent LLMs were called
    assert mock_llm.complete.call_count == 2
    assert sub_llm.complete.call_count == 2


@pytest.mark.asyncio
async def test_agent_multiple_tool_calls(mock_llm, calculator_tool):
    """Test agent making multiple tool calls in sequence.

    This simulates a more complex workflow:
    1. Calculate 10 + 5
    2. Multiply result by 2
    3. Return final result
    """
    # Setup mock LLM responses
    mock_llm.complete.side_effect = [
        # First calculation: 10 + 5
        LLMResponse(
            message=Message(
                role="assistant",
                content=None,
                tool_calls=[
                    ToolCall(
                        id="call-1",
                        name="calculator",
                        arguments={"operation": "add", "x": 10, "y": 5}
                    )
                ]
            ),
            finish_reason="tool_calls"
        ),
        # Second calculation: 15 * 2
        LLMResponse(
            message=Message(
                role="assistant",
                content=None,
                tool_calls=[
                    ToolCall(
                        id="call-2",
                        name="calculator",
                        arguments={"operation": "multiply", "x": 15, "y": 2}
                    )
                ]
            ),
            finish_reason="tool_calls"
        ),
        # Final response
        LLMResponse(
            message=Message(
                role="assistant",
                content="First I calculated 10 + 5 = 15, then 15 × 2 = 30."
            ),
            finish_reason="stop"
        )
    ]

    # Create agent
    agent = Agent(
        llm=mock_llm,
        tools=[calculator_tool],
        max_iterations=10
    )

    # Run agent
    result = await agent.run(
        user_message="Calculate (10 + 5) * 2",
        session_id="integration-test-3"
    )

    # Verify result
    assert isinstance(result, AgentResult)
    assert result.finish_reason == "stop"
    assert "30" in result.final_message.content

    # Verify message flow
    # user + assistant + tool + assistant + tool + assistant = 6
    assert len(result.messages) == 6
    assert result.messages[0].role == "user"
    assert result.messages[1].role == "assistant"
    assert result.messages[2].role == "tool"
    assert "Result: 15.0" in result.messages[2].content
    assert result.messages[3].role == "assistant"
    assert result.messages[4].role == "tool"
    assert "Result: 30.0" in result.messages[4].content
    assert result.messages[5].role == "assistant"

    # Verify LLM was called 3 times
    assert mock_llm.complete.call_count == 3
