"""Tests for specialized agent classes."""

import pytest
from unittest.mock import Mock

from agenai.agent import (
    AgentConfig,
    AgentContext,
    SpecializedAgent,
    ChainedAgent,
    RouterAgent,
)
from agenai.agent.tools import Tool, ToolResult
from agenai.agent.llm import LLMResponse, ToolCall


# =============================================================================
# Test Fixtures
# =============================================================================

class MockTool(Tool):
    """Mock tool for testing."""

    @property
    def name(self) -> str:
        return "mock_tool"

    @property
    def description(self) -> str:
        return "Mock tool"

    @property
    def parameters(self) -> dict:
        return {"type": "object", "properties": {}}

    def execute(self, **kwargs) -> ToolResult:
        return ToolResult.ok("mock result")


class MockLLMProvider:
    """Mock LLM for testing."""

    def __init__(self, response_text: str = "Mock response"):
        self.response_text = response_text
        self.call_count = 0

    def chat(self, messages, system=None, tools=None, max_tokens=4096, temperature=0.7):
        self.call_count += 1
        return LLMResponse(text=self.response_text, tool_calls=[])


# =============================================================================
# SpecializedAgent Tests
# =============================================================================

class TestSpecializedAgent:
    """Tests for SpecializedAgent base class."""

    def test_specialized_agent_definition(self):
        """Test defining a specialized agent."""

        class TestAgent(SpecializedAgent):
            name = "test"
            description = "Test agent"

            def get_system_prompt(self) -> str:
                return "You are a test agent."

            def get_tools(self) -> list:
                return [MockTool()]

        # Create with mock LLM
        llm = MockLLMProvider()
        agent = TestAgent(llm=llm)

        assert agent.name == "test"
        assert agent.description == "Test agent"
        assert "test agent" in agent.config.system_prompt.lower()
        assert len(agent.tools) == 1

    def test_specialized_agent_lifecycle_hooks(self):
        """Test specialized agent lifecycle hooks."""
        started = []
        completed = []

        class TestAgent(SpecializedAgent):
            name = "lifecycle_test"
            description = "Test lifecycle"

            def get_system_prompt(self) -> str:
                return "Test"

            def get_tools(self) -> list:
                return []

            def on_start(self, context):
                started.append(context.user_input)

            def on_complete(self, result):
                completed.append(result.output)

        llm = MockLLMProvider("Lifecycle complete")
        agent = TestAgent(llm=llm)

        result = agent.run("Test input")

        assert "Test input" in started
        assert "Lifecycle complete" in completed

    def test_specialized_agent_setup_called(self):
        """Test that setup() is called during init."""
        setup_called = []

        class TestAgent(SpecializedAgent):
            name = "setup_test"
            description = "Test setup"

            def get_system_prompt(self) -> str:
                return "Test"

            def get_tools(self) -> list:
                return []

            def setup(self):
                setup_called.append(True)

        llm = MockLLMProvider()
        TestAgent(llm=llm)

        assert len(setup_called) == 1

    def test_specialized_agent_custom_config(self):
        """Test specialized agent with custom config."""

        class TestAgent(SpecializedAgent):
            name = "config_test"
            description = "Test config"

            def get_system_prompt(self) -> str:
                return "Custom prompt"

            def get_tools(self) -> list:
                return []

        llm = MockLLMProvider()
        config = AgentConfig(max_steps=5, temperature=0.3)
        agent = TestAgent(config=config, llm=llm)

        assert agent.config.max_steps == 5
        assert agent.config.temperature == 0.3
        # System prompt should be overridden by get_system_prompt
        assert "Custom prompt" in agent.config.system_prompt


# =============================================================================
# ChainedAgent Tests
# =============================================================================

class TestChainedAgent:
    """Tests for ChainedAgent."""

    def test_chained_agent_sequential_execution(self):
        """Test that chained agents execute sequentially."""
        call_order = []

        class Agent1(SpecializedAgent):
            name = "agent1"
            description = "First"

            def get_system_prompt(self):
                return "First agent"

            def get_tools(self):
                return []

            def on_start(self, ctx):
                call_order.append(("agent1", ctx.user_input))

        class Agent2(SpecializedAgent):
            name = "agent2"
            description = "Second"

            def get_system_prompt(self):
                return "Second agent"

            def get_tools(self):
                return []

            def on_start(self, ctx):
                call_order.append(("agent2", ctx.user_input))

        llm1 = MockLLMProvider("Output from agent1")
        llm2 = MockLLMProvider("Final output")

        chain = ChainedAgent(
            agents=[
                Agent1(llm=llm1),
                Agent2(llm=llm2),
            ]
        )

        result = chain.run("Initial input")

        assert result.success
        assert "Final output" in result.output
        # Agent2 should receive agent1's output
        assert call_order[0][0] == "agent1"
        assert call_order[1][0] == "agent2"

    def test_chained_agent_failure_stops_chain(self):
        """Test that chain stops on agent failure."""

        class FailingAgent(SpecializedAgent):
            name = "failing"
            description = "Fails"

            def get_system_prompt(self):
                return "Fail"

            def get_tools(self):
                return []

        class SuccessAgent(SpecializedAgent):
            name = "success"
            description = "Success"
            was_called = False

            def get_system_prompt(self):
                return "Success"

            def get_tools(self):
                return []

            def on_start(self, ctx):
                SuccessAgent.was_called = True

        # Create a mock LLM that will cause max_steps failure
        class FailingLLM:
            def chat(self, *args, **kwargs):
                # Always return tool call to exhaust max_steps
                return LLMResponse(
                    tool_calls=[ToolCall(id="x", name="nonexistent", arguments={})]
                )

        failing_agent = FailingAgent(
            config=AgentConfig(max_steps=1),
            llm=FailingLLM()
        )

        chain = ChainedAgent(agents=[failing_agent])
        result = chain.run("Test")

        assert not result.success
        assert "failed" in result.error.lower() or "Max steps" in result.error


# =============================================================================
# RouterAgent Tests
# =============================================================================

class TestRouterAgent:
    """Tests for RouterAgent."""

    def test_router_routes_to_correct_agent(self):
        """Test that router selects correct agent."""

        class MathAgent(SpecializedAgent):
            name = "math"
            description = "Math problems"
            was_called = False

            def get_system_prompt(self):
                return "Math"

            def get_tools(self):
                return []

            def on_start(self, ctx):
                MathAgent.was_called = True

        class WriterAgent(SpecializedAgent):
            name = "writer"
            description = "Writing tasks"
            was_called = False

            def get_system_prompt(self):
                return "Writer"

            def get_tools(self):
                return []

            def on_start(self, ctx):
                WriterAgent.was_called = True

        # Router returns "math" for routing
        router_llm = MockLLMProvider("math")
        # Agents return their own responses
        math_llm = MockLLMProvider("Math result: 42")
        writer_llm = MockLLMProvider("Written content")

        router = RouterAgent(
            agents={
                "math": MathAgent(llm=math_llm),
                "writer": WriterAgent(llm=writer_llm),
            },
            llm=router_llm,
        )

        MathAgent.was_called = False
        WriterAgent.was_called = False

        result = router.run("What is 2+2?")

        assert result.success
        assert MathAgent.was_called
        assert not WriterAgent.was_called

    def test_router_fallback_to_first_agent(self):
        """Test router falls back when no match found."""

        class DefaultAgent(SpecializedAgent):
            name = "default"
            description = "Default"

            def get_system_prompt(self):
                return "Default"

            def get_tools(self):
                return []

        # Router returns nonsense
        router_llm = MockLLMProvider("nonexistent_agent")
        default_llm = MockLLMProvider("Default response")

        router = RouterAgent(
            agents={"default": DefaultAgent(llm=default_llm)},
            llm=router_llm,
        )

        result = router.run("Something")

        assert result.success
        assert "Default response" in result.output


# =============================================================================
# Edge Cases and Error Handling
# =============================================================================

class TestEdgeCases:
    """Tests for edge cases and error handling."""

    def test_agent_with_no_tools(self):
        """Test agent works without tools."""

        class NoToolsAgent(SpecializedAgent):
            name = "no_tools"
            description = "No tools"

            def get_system_prompt(self):
                return "No tools needed"

            def get_tools(self):
                return []

        llm = MockLLMProvider("Response without tools")
        agent = NoToolsAgent(llm=llm)

        result = agent.run("Hello")

        assert result.success
        assert len(agent.tools) == 0

    def test_agent_handles_empty_input(self):
        """Test agent handles empty input."""

        class TestAgent(SpecializedAgent):
            name = "test"
            description = "Test"

            def get_system_prompt(self):
                return "Test"

            def get_tools(self):
                return []

        llm = MockLLMProvider("Handled empty")
        agent = TestAgent(llm=llm)

        result = agent.run("")

        assert result.success

    def test_multiple_tool_calls_in_single_step(self):
        """Test handling multiple tool calls in one LLM response."""

        class Tool1(Tool):
            @property
            def name(self):
                return "tool1"

            @property
            def description(self):
                return "Tool 1"

            @property
            def parameters(self):
                return {"type": "object", "properties": {}}

            def execute(self, **kwargs):
                Tool1.call_count = getattr(Tool1, 'call_count', 0) + 1
                return ToolResult.ok("ok")

        class Tool2(Tool):
            @property
            def name(self):
                return "tool2"

            @property
            def description(self):
                return "Tool 2"

            @property
            def parameters(self):
                return {"type": "object", "properties": {}}

            def execute(self, **kwargs):
                Tool2.call_count = getattr(Tool2, 'call_count', 0) + 1
                return ToolResult.ok("ok")

        tool1 = Tool1()
        tool2 = Tool2()
        Tool1.call_count = 0
        Tool2.call_count = 0

        class MultiToolLLM:
            def __init__(self):
                self.call_count = 0

            def chat(self, *args, **kwargs):
                self.call_count += 1
                if self.call_count == 1:
                    return LLMResponse(
                        text="Calling multiple tools",
                        tool_calls=[
                            ToolCall(id="1", name="tool1", arguments={}),
                            ToolCall(id="2", name="tool2", arguments={}),
                        ]
                    )
                return LLMResponse(text="Done", tool_calls=[])

        from agenai.agent import Agent
        from agenai.agent.tools import ToolRegistry

        tools = ToolRegistry()
        tools.register(tool1)
        tools.register(tool2)

        agent = Agent(llm=MultiToolLLM(), tools=tools)
        result = agent.run("Call both tools")

        assert result.success
        assert Tool1.call_count == 1
        assert Tool2.call_count == 1
