"""Tests for core agent components."""

import pytest
from unittest.mock import Mock, MagicMock, patch
from dataclasses import dataclass

from agenai.agent import (
    Agent,
    AgentConfig,
    AgentContext,
    AgentResult,
    StepResult,
    StepType,
    ConversationMemory,
    HookRegistry,
    HookEvent,
    HookData,
    RetryConfig,
    with_retry,
    RetryError,
)
from agenai.agent.tools import Tool, ToolRegistry, ToolResult
from agenai.agent.llm import LLMProvider, LLMResponse, ToolCall


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

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

    def __init__(self, name: str = "mock_tool", return_value: str = "mock result"):
        self._name = name
        self._return_value = return_value
        self.call_count = 0
        self.last_args = None

    @property
    def name(self) -> str:
        return self._name

    @property
    def description(self) -> str:
        return "A mock tool for testing"

    @property
    def parameters(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "input": {"type": "string", "description": "Input value"}
            },
            "required": ["input"],
        }

    def execute(self, **kwargs) -> ToolResult:
        self.call_count += 1
        self.last_args = kwargs
        return ToolResult.ok(self._return_value)


class MockLLMProvider(LLMProvider):
    """Mock LLM provider for testing."""

    def __init__(self, responses: list[LLMResponse] | None = None):
        self.responses = responses or []
        self.call_count = 0
        self.last_messages = None

    def chat(self, messages, system=None, tools=None, max_tokens=4096, temperature=0.7):
        self.call_count += 1
        self.last_messages = messages

        if self.call_count <= len(self.responses):
            return self.responses[self.call_count - 1]

        # Default: return final answer
        return LLMResponse(text="Default response", tool_calls=[])


# =============================================================================
# AgentConfig Tests
# =============================================================================

class TestAgentConfig:
    """Tests for AgentConfig."""

    def test_default_config(self):
        """Test default configuration values."""
        config = AgentConfig()

        assert config.model == "claude-sonnet-4-6"
        assert config.max_steps == 20
        assert config.max_tokens == 4096
        assert 0 <= config.temperature <= 1

    def test_custom_config(self):
        """Test custom configuration."""
        config = AgentConfig(
            model="claude-opus-4-6",
            max_steps=10,
            temperature=0.5,
            system_prompt="Custom prompt",
        )

        assert config.model == "claude-opus-4-6"
        assert config.max_steps == 10
        assert config.temperature == 0.5
        assert config.system_prompt == "Custom prompt"

    def test_config_validation(self):
        """Test configuration validation."""
        # Temperature must be 0-1
        with pytest.raises(ValueError):
            AgentConfig(temperature=1.5)

        # Max steps must be positive
        with pytest.raises(ValueError):
            AgentConfig(max_steps=0)


# =============================================================================
# ToolRegistry Tests
# =============================================================================

class TestToolRegistry:
    """Tests for ToolRegistry."""

    def test_register_tool(self):
        """Test registering a tool."""
        registry = ToolRegistry()
        tool = MockTool("test_tool")

        registry.register(tool)

        assert "test_tool" in registry
        assert len(registry) == 1

    def test_duplicate_registration_fails(self):
        """Test that registering duplicate tool fails."""
        registry = ToolRegistry()
        tool = MockTool("test_tool")

        registry.register(tool)

        with pytest.raises(ValueError, match="already registered"):
            registry.register(tool)

    def test_get_tool(self):
        """Test getting a tool by name."""
        registry = ToolRegistry()
        tool = MockTool("test_tool")
        registry.register(tool)

        retrieved = registry.get("test_tool")

        assert retrieved is tool

    def test_get_missing_tool_fails(self):
        """Test that getting missing tool raises KeyError."""
        registry = ToolRegistry()

        with pytest.raises(KeyError):
            registry.get("nonexistent")

    def test_execute_tool(self):
        """Test executing a tool through registry."""
        registry = ToolRegistry()
        tool = MockTool("test_tool", return_value="executed!")
        registry.register(tool)

        result = registry.execute("test_tool", input="test")

        assert result.success
        assert result.output == "executed!"
        assert tool.call_count == 1

    def test_list_tools(self):
        """Test listing tool schemas."""
        registry = ToolRegistry()
        registry.register(MockTool("tool1"))
        registry.register(MockTool("tool2"))

        schemas = registry.list()

        assert len(schemas) == 2
        names = [s["name"] for s in schemas]
        assert "tool1" in names
        assert "tool2" in names


# =============================================================================
# ConversationMemory Tests
# =============================================================================

class TestConversationMemory:
    """Tests for ConversationMemory."""

    def test_add_user_message(self):
        """Test adding user messages."""
        memory = ConversationMemory()

        memory.add_user("Hello")

        assert len(memory) == 1
        messages = memory.get_messages()
        assert messages[0].content == "Hello"

    def test_add_assistant_message(self):
        """Test adding assistant messages."""
        memory = ConversationMemory()

        memory.add_assistant("Hi there!")

        assert len(memory) == 1

    def test_add_tool_result(self):
        """Test adding tool results."""
        memory = ConversationMemory()

        memory.add_tool_result("tool_123", "my_tool", "result data")

        assert len(memory) == 1
        msg = memory.get_messages()[0]
        assert msg.tool_use_id == "tool_123"

    def test_clear_memory(self):
        """Test clearing memory."""
        memory = ConversationMemory()
        memory.add_user("msg1")
        memory.add_user("msg2")

        memory.clear()

        assert len(memory) == 0

    def test_truncate_memory(self):
        """Test truncating to max messages."""
        memory = ConversationMemory()
        for i in range(10):
            memory.add_user(f"message {i}")

        memory.truncate(5)

        assert len(memory) == 5
        # Should keep most recent
        assert "message 9" in memory.get_messages()[-1].content


# =============================================================================
# HookRegistry Tests
# =============================================================================

class TestHookRegistry:
    """Tests for HookRegistry."""

    def test_register_and_trigger_hook(self):
        """Test registering and triggering hooks."""
        hooks = HookRegistry()
        triggered = []

        @hooks.on(HookEvent.TOOL_START)
        def on_tool(data):
            triggered.append(data.event)

        hooks.trigger(HookData(event=HookEvent.TOOL_START))

        assert HookEvent.TOOL_START in triggered

    def test_multiple_events(self):
        """Test hook listening to multiple events."""
        hooks = HookRegistry()
        triggered = []

        @hooks.on(HookEvent.TOOL_START, HookEvent.TOOL_END)
        def on_tool(data):
            triggered.append(data.event)

        hooks.trigger(HookData(event=HookEvent.TOOL_START))
        hooks.trigger(HookData(event=HookEvent.TOOL_END))

        assert len(triggered) == 2

    def test_hook_receives_data(self):
        """Test that hooks receive correct data."""
        hooks = HookRegistry()
        received_data = []

        @hooks.on(HookEvent.TOOL_START)
        def on_tool(data):
            received_data.append(data)

        hooks.trigger(HookData(
            event=HookEvent.TOOL_START,
            tool_name="test_tool",
            tool_args={"key": "value"},
        ))

        assert len(received_data) == 1
        assert received_data[0].tool_name == "test_tool"
        assert received_data[0].tool_args == {"key": "value"}

    def test_hook_error_doesnt_crash(self):
        """Test that hook errors don't crash the system."""
        hooks = HookRegistry()

        @hooks.on(HookEvent.TOOL_START)
        def bad_hook(data):
            raise ValueError("Hook error!")

        # Should not raise
        hooks.trigger(HookData(event=HookEvent.TOOL_START))


# =============================================================================
# Retry Tests
# =============================================================================

class TestRetry:
    """Tests for retry logic."""

    def test_successful_on_first_try(self):
        """Test function succeeds on first try."""
        def success():
            return "ok"

        result = with_retry(success, RetryConfig(max_retries=3, base_delay=0.001))

        assert result == "ok"

    def test_retry_on_failure(self):
        """Test retry after failure."""
        attempts = [0]

        def flaky():
            attempts[0] += 1
            if attempts[0] < 3:
                raise ValueError("fail")
            return "success"

        result = with_retry(flaky, RetryConfig(max_retries=3, base_delay=0.001))

        assert result == "success"
        assert attempts[0] == 3

    def test_max_retries_exceeded(self):
        """Test failure when max retries exceeded."""
        def always_fail():
            raise ValueError("always fails")

        with pytest.raises(RetryError) as exc_info:
            with_retry(always_fail, RetryConfig(max_retries=2, base_delay=0.001))

        assert exc_info.value.attempts == 3  # initial + 2 retries


# =============================================================================
# AgentContext Tests
# =============================================================================

class TestAgentContext:
    """Tests for AgentContext."""

    def test_state_management(self):
        """Test getting and setting state."""
        ctx = AgentContext()

        ctx.set("key", "value")

        assert ctx.get("key") == "value"
        assert ctx.get("missing", "default") == "default"

    def test_child_context(self):
        """Test creating child context."""
        parent = AgentContext(run_id="parent", state={"shared": "data"})

        child = parent.child(run_id="child")

        assert child.parent is parent
        assert child.get("shared") == "data"
        assert child.run_id == "child"

    def test_record_tool_call(self):
        """Test recording tool calls."""
        ctx = AgentContext()

        ctx.record_tool_call("my_tool", {"arg": "val"}, "result", True)

        assert len(ctx.tool_calls) == 1
        assert ctx.tool_calls[0]["name"] == "my_tool"


# =============================================================================
# Agent Integration Tests
# =============================================================================

class TestAgentIntegration:
    """Integration tests for Agent."""

    def test_agent_with_final_answer(self):
        """Test agent that returns final answer immediately."""
        llm = MockLLMProvider([
            LLMResponse(text="Hello, I'm here to help!", tool_calls=[])
        ])

        agent = Agent(
            config=AgentConfig(system_prompt="Be helpful"),
            llm=llm,
        )

        result = agent.run("Hi")

        assert result.success
        assert "Hello" in result.output
        assert len(result.steps) == 1
        assert result.steps[0].step_type == StepType.FINAL_ANSWER

    def test_agent_with_tool_call(self):
        """Test agent that calls a tool then answers."""
        tool = MockTool("get_data", return_value="tool data")
        tools = ToolRegistry()
        tools.register(tool)

        llm = MockLLMProvider([
            # First response: call tool
            LLMResponse(
                text="Let me check that.",
                tool_calls=[ToolCall(id="call_1", name="get_data", arguments={"input": "test"})]
            ),
            # Second response: final answer
            LLMResponse(text="Based on the data: tool data", tool_calls=[])
        ])

        agent = Agent(llm=llm, tools=tools)
        result = agent.run("Get me some data")

        assert result.success
        assert tool.call_count == 1
        assert len(result.steps) == 2
        assert result.steps[0].step_type == StepType.TOOL_CALL
        assert result.steps[1].step_type == StepType.FINAL_ANSWER

    def test_agent_max_steps(self):
        """Test agent stops at max steps."""
        llm = MockLLMProvider([
            # Always call a tool, never finish
            LLMResponse(
                tool_calls=[ToolCall(id=f"call_{i}", name="mock_tool", arguments={"input": "x"})]
            )
            for i in range(10)
        ])

        tools = ToolRegistry()
        tools.register(MockTool())

        agent = Agent(
            config=AgentConfig(max_steps=3),
            llm=llm,
            tools=tools,
        )

        result = agent.run("Loop forever")

        assert not result.success
        assert "Max steps" in result.error
        assert len(result.steps) == 3

    def test_agent_hooks_triggered(self):
        """Test that agent triggers hooks."""
        events = []
        hooks = HookRegistry()

        @hooks.on(HookEvent.AGENT_START, HookEvent.AGENT_END, HookEvent.STEP_START, HookEvent.STEP_END)
        def track(data):
            events.append(data.event)

        llm = MockLLMProvider([LLMResponse(text="Done", tool_calls=[])])
        agent = Agent(llm=llm, hooks=hooks)

        agent.run("Test")

        assert HookEvent.AGENT_START in events
        assert HookEvent.AGENT_END in events
        assert HookEvent.STEP_START in events
        assert HookEvent.STEP_END in events

    def test_agent_context_passed(self):
        """Test that context is properly passed."""
        llm = MockLLMProvider([LLMResponse(text="Done", tool_calls=[])])
        agent = Agent(llm=llm)

        result = agent.run("Test")

        assert result.context is not None
        assert result.context.user_input == "Test"
        assert result.context.run_id != ""

    def test_agent_reset(self):
        """Test agent reset clears state."""
        llm = MockLLMProvider([
            LLMResponse(text="First", tool_calls=[]),
            LLMResponse(text="Second", tool_calls=[]),
        ])
        agent = Agent(llm=llm)

        agent.run("First message")
        assert len(agent.memory) > 0

        agent.reset()

        assert len(agent.memory) == 0
        assert agent.context is None


# =============================================================================
# ToolResult Tests
# =============================================================================

class TestToolResult:
    """Tests for ToolResult."""

    def test_ok_result(self):
        """Test creating successful result."""
        result = ToolResult.ok("success output")

        assert result.success
        assert result.output == "success output"
        assert result.error is None

    def test_fail_result(self):
        """Test creating failed result."""
        result = ToolResult.fail("error message")

        assert not result.success
        assert result.output == ""
        assert result.error == "error message"
