"""Tests for Agent execution loop."""

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

from yier_agents.src.agent import Agent, AgentResult
from yier_agents.src.memory import CompactionConfig, JSONSessionStore
from yier_agents.src.message import Message, ToolCall, LLMResponse
from yier_agents.src.skill import Skill, SkillCatalog, SkillDocument
from yier_agents.src.tool import Tool, ToolContext, ToolOutput
from yier_agents.src.llm import LLM


# Define test tool parameters
class EchoParams(BaseModel):
    """Parameters for echo tool."""
    text: str = Field(..., description="Text to echo")


async def echo_execute(params: EchoParams, ctx: ToolContext) -> ToolOutput:
    """Execute echo operation."""
    return ToolOutput(content=f"Echo: {params.text}")


# Test fixtures
@pytest.fixture
def echo_tool():
    """Create an echo tool."""
    return Tool(
        name="echo",
        description="Echo the input text",
        parameters=EchoParams,
        execute=echo_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_agent_simple_conversation(mock_llm):
    """Test simple conversation without tool calls."""
    # Setup mock LLM to return a simple response
    mock_llm.complete.return_value = LLMResponse(
        message=Message(role="assistant", content="Hello! How can I help you?"),
        finish_reason="stop",
        usage=18  # Total tokens
    )

    # Create agent without tools
    agent = Agent(llm=mock_llm, tools=[], max_iterations=20)

    # Run agent
    result = await agent.run("Hello", session_id="test-session-1")

    # Assertions
    assert isinstance(result, AgentResult)
    assert result.session_id == "test-session-1"
    assert result.finish_reason == "stop"
    assert result.final_message.role == "assistant"
    assert result.final_message.content == "Hello! How can I help you?"
    assert len(result.messages) == 2  # user + assistant
    assert result.messages[0].role == "user"
    assert result.messages[0].content == "Hello"

    # Verify LLM was called once
    assert mock_llm.complete.call_count == 1


@pytest.mark.asyncio
async def test_agent_with_tool_call(mock_llm, echo_tool):
    """Test agent with tool call flow."""
    # Setup mock LLM to return tool call first, then final response
    mock_llm.complete.side_effect = [
        # First call: LLM decides to use echo tool
        LLMResponse(
            message=Message(
                role="assistant",
                content=None,
                tool_calls=[
                    ToolCall(
                        id="call-1",
                        name="echo",
                        arguments={"text": "test message"}
                    )
                ]
            ),
            finish_reason="tool_calls"
        ),
        # Second call: LLM responds with final answer
        LLMResponse(
            message=Message(
                role="assistant",
                content="I echoed your message successfully."
            ),
            finish_reason="stop"
        )
    ]

    # Create agent with echo tool
    agent = Agent(llm=mock_llm, tools=[echo_tool], max_iterations=20)

    # Run agent
    result = await agent.run("Echo 'test message'", session_id="test-session-2")

    # Assertions
    assert isinstance(result, AgentResult)
    assert result.session_id == "test-session-2"
    assert result.finish_reason == "stop"
    assert result.final_message.content == "I echoed your message successfully."

    # Check message flow: user -> assistant (tool call) -> tool -> assistant (final)
    assert len(result.messages) == 4
    assert result.messages[0].role == "user"
    assert result.messages[1].role == "assistant"
    assert result.messages[1].tool_calls is not None
    assert result.messages[2].role == "tool"
    assert result.messages[2].content == "Echo: test message"
    assert result.messages[3].role == "assistant"

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


@pytest.mark.asyncio
async def test_agent_max_iterations(mock_llm, echo_tool):
    """Test max iterations protection."""
    # Setup mock LLM to always return tool calls (infinite loop scenario)
    mock_llm.complete.return_value = LLMResponse(
        message=Message(
            role="assistant",
            content=None,
            tool_calls=[
                ToolCall(
                    id="call-loop",
                    name="echo",
                    arguments={"text": "loop"}
                )
            ]
        ),
        finish_reason="tool_calls"
    )

    # Create agent with max_iterations=3
    agent = Agent(llm=mock_llm, tools=[echo_tool], max_iterations=3)

    # Run agent
    result = await agent.run("Start loop", session_id="test-session-3")

    # Assertions
    assert isinstance(result, AgentResult)
    assert result.session_id == "test-session-3"
    assert result.finish_reason == "max_iterations"

    # Should have stopped at max iterations
    # Each iteration: assistant (tool call) + tool result = 2 messages
    # Plus initial user message = 1
    # 3 iterations: 1 + (2 * 3) = 7 messages
    assert len(result.messages) == 7

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


@pytest.mark.asyncio
async def test_agent_system_prompt(mock_llm):
    """Test system prompt is added as first message."""
    # Setup mock LLM
    mock_llm.complete.return_value = LLMResponse(
        message=Message(role="assistant", content="I am a helpful assistant."),
        finish_reason="stop"
    )

    # Create agent with system prompt
    system_prompt = "You are a helpful assistant that always responds politely."
    agent = Agent(
        llm=mock_llm,
        tools=[],
        max_iterations=20,
        system_prompt=system_prompt
    )

    # Run agent
    result = await agent.run("Hello", session_id="test-session-4")

    # Assertions
    assert isinstance(result, AgentResult)
    assert len(result.messages) == 3  # system + user + assistant
    assert result.messages[0].role == "system"
    assert result.messages[0].content == system_prompt
    assert result.messages[1].role == "user"
    assert result.messages[1].content == "Hello"
    assert result.messages[2].role == "assistant"

    # Verify LLM received system prompt
    call_args = mock_llm.complete.call_args
    messages_passed = call_args[1]["messages"]
    assert messages_passed[0].role == "system"
    assert messages_passed[0].content == system_prompt


@pytest.mark.asyncio
async def test_agent_merges_skill_tools_and_prompts(mock_llm, echo_tool):
    """Test skill-provided tools and prompts are applied to the agent."""
    mock_llm.complete.return_value = LLMResponse(
        message=Message(role="assistant", content="Skills loaded."),
        finish_reason="stop",
    )

    skill = Skill(
        name="echo-skill",
        description="Adds echo support",
        tools=[echo_tool],
        system_prompt="Use the echo tool when repetition helps.",
    )
    agent = Agent(
        llm=mock_llm,
        tools=[],
        skills=[skill],
        system_prompt="You are a helpful assistant.",
    )

    result = await agent.run("Say hello", session_id="skill-session")

    assert result.finish_reason == "stop"
    call_args = mock_llm.complete.call_args
    messages_passed = call_args.kwargs["messages"]
    assert messages_passed[0].role == "system"
    assert "[Skill: echo-skill]" in messages_passed[0].content
    assert "Use the echo tool when repetition helps." in messages_passed[0].content

    tools_passed = call_args.kwargs["tools"]
    assert len(tools_passed) == 1
    assert tools_passed[0]["function"]["name"] == "echo"


@pytest.mark.asyncio
async def test_agent_uses_json_session_store(mock_llm, tmp_path):
    """Test session memory can be persisted to disk."""
    mock_llm.complete.side_effect = [
        LLMResponse(
            message=Message(role="assistant", content="Stored."),
            finish_reason="stop",
        ),
        LLMResponse(
            message=Message(role="assistant", content="Remembered."),
            finish_reason="stop",
        ),
    ]

    session_store = JSONSessionStore(tmp_path / "sessions")
    first_agent = Agent(
        llm=mock_llm,
        tools=[],
        session_store=session_store,
        system_prompt="Persistent assistant",
    )

    await first_agent.run("Remember this", session_id="persisted-session")

    second_agent = Agent(
        llm=mock_llm,
        tools=[],
        session_store=session_store,
        system_prompt="Persistent assistant",
    )
    await second_agent.run("What did I say?", session_id="persisted-session")

    second_call_args = mock_llm.complete.call_args_list[1]
    restored_messages = second_call_args.kwargs["messages"]
    assert [message.role for message in restored_messages[:4]] == [
        "system",
        "user",
        "assistant",
        "user",
    ]
    assert restored_messages[1].content == "Remember this"
    assert restored_messages[2].content == "Stored."


@pytest.mark.asyncio
async def test_agent_compacts_history_into_summary(mock_llm):
    """Test older history is summarized when compaction is enabled."""
    mock_llm.complete.side_effect = [
        LLMResponse(
            message=Message(role="assistant", content="Summary of previous messages."),
            finish_reason="stop",
        ),
        LLMResponse(
            message=Message(role="assistant", content="Final response."),
            finish_reason="stop",
        ),
    ]

    agent = Agent(
        llm=mock_llm,
        tools=[],
        system_prompt="Helpful assistant",
        compaction_config=CompactionConfig(
            enabled=True,
            trigger_message_count=6,
            preserve_recent_messages=2,
            summary_max_tokens=128,
        ),
    )
    session_id = "compact-session"
    history = [
        Message(role="system", content="Helpful assistant"),
        Message(role="user", content="Message 1"),
        Message(role="assistant", content="Reply 1"),
        Message(role="user", content="Message 2"),
        Message(role="assistant", content="Reply 2"),
        Message(role="user", content="Message 3"),
    ]
    agent.session_store.save(session_id, history)

    await agent.run("Message 4", session_id=session_id)

    summarizer_call = mock_llm.complete.call_args_list[0]
    assert summarizer_call.kwargs["tools"] is None

    final_call = mock_llm.complete.call_args_list[1]
    compacted_messages = final_call.kwargs["messages"]
    assert compacted_messages[0].role == "system"
    assert compacted_messages[0].content == "Helpful assistant"
    assert compacted_messages[1].content.startswith("<conversation_summary>")
    assert "Summary of previous messages." in compacted_messages[1].content
    assert compacted_messages[-2].content == "Message 4"


@pytest.mark.asyncio
async def test_agent_exposes_lazy_skill_catalog(mock_llm, tmp_path):
    """Test Agent injects available_skills and the skill_load tool."""
    mock_llm.complete.return_value = LLMResponse(
        message=Message(role="assistant", content="Loaded."),
        finish_reason="stop",
    )

    skill_dir = tmp_path / ".agents" / "skills" / "docs-skill"
    skill_dir.mkdir(parents=True)
    skill_file = skill_dir / "SKILL.md"
    skill_file.write_text(
        "\n".join(
            [
                "---",
                "name: docs-skill",
                "description: Helps with documentation tasks.",
                "---",
                "",
                "Use this skill for docs work.",
            ]
        ),
        encoding="utf-8",
    )

    catalog = SkillCatalog(
        [
            SkillDocument(
                name="docs-skill",
                description="Helps with documentation tasks.",
                location=skill_file.resolve(),
                content="Use this skill for docs work.",
            )
        ]
    )
    agent = Agent(
        llm=mock_llm,
        tools=[],
        skill_catalog=catalog,
        system_prompt="You are helpful.",
    )

    await agent.run("Need docs help", session_id="skill-catalog-session")

    call_args = mock_llm.complete.call_args
    messages_passed = call_args.kwargs["messages"]
    assert "<available_skills>" in messages_passed[0].content
    assert "<name>docs-skill</name>" in messages_passed[0].content

    tools_passed = call_args.kwargs["tools"]
    assert any(tool_schema["function"]["name"] == "skill_load" for tool_schema in tools_passed)
