"""Tests for the ReAct agent loop (app/ai/agent_loop.py).

All tests mock llm_complete and tool_registry — no real LLM or DB needed.
"""

from __future__ import annotations

import asyncio
import json
import uuid
from dataclasses import dataclass
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from app.ai.agent_loop import ReActResult, ReActStep, run_react_loop


# ──────────────────────────────────────────────────────────────────────────
# Test helpers
# ──────────────────────────────────────────────────────────────────────────


@dataclass
class MockAgentDefinition:
    """Minimal stand-in for AgentDefinition used by run_react_loop."""

    id: uuid.UUID
    llm_model: str = "gpt-4o"
    system_prompt: str = "You are a helpful assistant."
    api_key: str | None = None
    api_base: str | None = None
    provider: str | None = None
    max_tokens: int = 1000


def _make_tool_call(
    call_id: str = "call_1",
    name: str = "search",
    arguments: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Build a mock LiteLLM tool-call object."""
    args_str = json.dumps(arguments or {})
    function_mock = MagicMock()
    function_mock.name = name
    function_mock.arguments = args_str
    tc = MagicMock()
    tc.id = call_id
    tc.function = function_mock
    return tc


def _make_llm_response(
    content: str = "",
    tool_calls: list[Any] | None = None,
    cost_usd: float = 0.001,
) -> dict[str, Any]:
    """Build a mock llm_complete() return value."""
    message = MagicMock()
    message.content = content
    message.tool_calls = tool_calls or None

    choice = MagicMock()
    choice.message = message

    raw_response = MagicMock()
    raw_response.choices = [choice]

    return {
        "content": content,
        "usage": {"total_tokens": 100},
        "cost_usd": cost_usd,
        "model": "gpt-4o",
        "raw_response": raw_response,
    }


def _make_tool_registry(tools: dict[str, AsyncMock] | None = None) -> MagicMock:
    """Build a mock ToolRegistry."""
    registry = MagicMock()
    _tools = tools or {}

    def _get(name: str) -> Any:
        if name in _tools:
            tool = MagicMock()
            tool.handler = _tools[name]
            return tool
        return None

    registry.get = _get
    return registry


@pytest.fixture
def agent_def() -> MockAgentDefinition:
    return MockAgentDefinition(id=uuid.uuid4())


@pytest.fixture
def tenant_id() -> uuid.UUID:
    return uuid.uuid4()


@pytest.fixture
def user_id() -> uuid.UUID:
    return uuid.uuid4()


@pytest.fixture
def mock_db() -> AsyncMock:
    return AsyncMock()


@pytest.fixture(autouse=True)
def _mock_hooks():
    """Patch do_action so the loop doesn't try to fire real hooks."""
    with patch("app.core.hooks.do_action", new_callable=AsyncMock):
        yield


# ──────────────────────────────────────────────────────────────────────────
# Tests
# ──────────────────────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_basic_react_loop_no_tools(agent_def, tenant_id, user_id, mock_db):
    """Single-step loop: LLM returns final response, no tool calls."""
    registry = _make_tool_registry()

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.return_value = _make_llm_response(
            content="Hello! How can I help you?",
            cost_usd=0.002,
        )

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Hi"}],
            tools=[],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
        )

    assert result.status == "completed"
    assert result.final_content == "Hello! How can I help you?"
    assert result.steps_taken == 1
    assert len(result.steps) == 1
    assert result.steps[0].action is None  # No tool call
    assert result.steps[0].thought == "Hello! How can I help you?"
    assert result.total_cost_usd == pytest.approx(0.002)
    mock_llm.assert_awaited_once()


@pytest.mark.asyncio
async def test_multi_step_loop_with_tool_calls(agent_def, tenant_id, user_id, mock_db):
    """Multi-step loop: LLM calls a tool, gets result, then gives final answer."""
    search_handler = AsyncMock(return_value="Found 3 contacts: Alice, Bob, Charlie")
    registry = _make_tool_registry({"search": search_handler})

    # Step 1: LLM calls search tool
    # Step 2: LLM gives final answer
    responses = [
        _make_llm_response(
            content="I'll search for contacts.",
            tool_calls=[_make_tool_call(name="search", arguments={"query": "contacts"})],
            cost_usd=0.003,
        ),
        _make_llm_response(
            content="I found 3 contacts: Alice, Bob, and Charlie.",
            cost_usd=0.004,
        ),
    ]

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = responses

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Find contacts"}],
            tools=[{"type": "function", "function": {"name": "search"}}],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
        )

    assert result.status == "completed"
    assert result.steps_taken == 2
    assert len(result.steps) == 2
    assert result.steps[0].action == "search"
    assert result.steps[0].action_input == {"query": "contacts"}
    assert result.steps[0].observation == "Found 3 contacts: Alice, Bob, Charlie"
    assert result.steps[1].action is None  # Final response
    assert result.final_content == "I found 3 contacts: Alice, Bob, and Charlie."
    assert result.total_cost_usd == pytest.approx(0.007)
    assert mock_llm.await_count == 2
    search_handler.assert_awaited_once()


@pytest.mark.asyncio
async def test_max_steps_limit(agent_def, tenant_id, user_id, mock_db):
    """Loop stops gracefully when max_steps is reached."""
    search_handler = AsyncMock(return_value="result")
    registry = _make_tool_registry({"search": search_handler})

    # Every response has a tool call — never gives final answer
    tool_call = _make_tool_call(name="search", arguments={})
    response = _make_llm_response(
        content="Searching...",
        tool_calls=[tool_call],
        cost_usd=0.001,
    )

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.return_value = response

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Keep searching"}],
            tools=[{"type": "function", "function": {"name": "search"}}],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
            max_steps=3,
        )

    assert result.status == "stopped_max_steps"
    assert result.steps_taken == 3
    assert len(result.steps) == 3
    assert "max_steps" in (result.error or "")
    assert mock_llm.await_count == 3


@pytest.mark.asyncio
async def test_timeout_graceful_stop(agent_def, tenant_id, user_id, mock_db):
    """Loop stops gracefully when timeout is exceeded."""
    registry = _make_tool_registry()

    # Simulate slow LLM responses that eventually exceed timeout
    async def slow_llm(**kwargs: Any) -> dict[str, Any]:
        await asyncio.sleep(0.5)
        return _make_llm_response(content="thinking...", tool_calls=None, cost_usd=0.001)

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = slow_llm

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Hi"}],
            tools=[],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
            timeout_seconds=0,  # Immediate timeout on first check
        )

    assert result.status == "stopped_timeout"
    assert "timeout" in (result.error or "").lower()


@pytest.mark.asyncio
async def test_transient_error_retry(agent_def, tenant_id, user_id, mock_db):
    """Transient errors are retried, then succeed."""
    registry = _make_tool_registry()

    # First call raises transient error, second succeeds
    responses: list[Any] = [
        Exception("rate limit exceeded"),
        _make_llm_response(content="Success after retry", cost_usd=0.002),
    ]

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = responses
        with patch("app.ai.agent_loop.asyncio.sleep", new_callable=AsyncMock):
            result = await run_react_loop(
                agent_definition=agent_def,
                messages=[{"role": "user", "content": "Hi"}],
                tools=[],
                tool_registry=registry,
                db=mock_db,
                tenant_id=tenant_id,
                user_id=user_id,
            )

    assert result.status == "completed"
    assert result.final_content == "Success after retry"
    assert mock_llm.await_count == 2  # First failed, second succeeded


@pytest.mark.asyncio
async def test_permanent_error_stops(agent_def, tenant_id, user_id, mock_db):
    """Permanent errors stop the loop immediately."""
    registry = _make_tool_registry()

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = Exception("authentication error: invalid api key")

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Hi"}],
            tools=[],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
        )

    assert result.status == "stopped_error"
    assert "Permanent error" in (result.error or "")
    assert mock_llm.await_count == 1  # No retry for permanent errors


@pytest.mark.asyncio
async def test_cost_accumulation(agent_def, tenant_id, user_id, mock_db):
    """Cost is accumulated across multiple LLM calls."""
    search_handler = AsyncMock(return_value="result")
    registry = _make_tool_registry({"search": search_handler})

    responses = [
        _make_llm_response(
            content="Step 1",
            tool_calls=[_make_tool_call(name="search")],
            cost_usd=0.01,
        ),
        _make_llm_response(
            content="Step 2",
            tool_calls=[_make_tool_call(name="search")],
            cost_usd=0.02,
        ),
        _make_llm_response(
            content="Final answer",
            cost_usd=0.03,
        ),
    ]

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = responses

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Search twice"}],
            tools=[{"type": "function", "function": {"name": "search"}}],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
        )

    assert result.status == "completed"
    assert result.total_cost_usd == pytest.approx(0.06)  # 0.01 + 0.02 + 0.03
    assert result.steps_taken == 3


@pytest.mark.asyncio
async def test_step_persistence_via_callback(agent_def, tenant_id, user_id, mock_db):
    """Steps are passed to on_step callback for persistence."""
    registry = _make_tool_registry()
    collected_steps: list[ReActStep] = []

    async def on_step(step: ReActStep) -> None:
        collected_steps.append(step)

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.return_value = _make_llm_response(content="Done", cost_usd=0.001)

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Hi"}],
            tools=[],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
            on_step=on_step,
        )

    assert len(collected_steps) == 1
    assert collected_steps[0].step_number == 1
    assert collected_steps[0].thought == "Done"
    assert result.status == "completed"


@pytest.mark.asyncio
async def test_tool_not_found(agent_def, tenant_id, user_id, mock_db):
    """When a tool is not found, the observation contains an error message."""
    registry = _make_tool_registry()  # No tools registered

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = [
            _make_llm_response(
                content="Calling unknown tool",
                tool_calls=[_make_tool_call(name="nonexistent", arguments={})],
                cost_usd=0.001,
            ),
            _make_llm_response(content="OK", cost_usd=0.001),
        ]

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Call unknown tool"}],
            tools=[{"type": "function", "function": {"name": "nonexistent"}}],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
        )

    assert result.status == "completed"
    assert result.steps[0].action == "nonexistent"
    assert "not found" in (result.steps[0].observation or "").lower()


@pytest.mark.asyncio
async def test_multiple_tool_calls_per_step(agent_def, tenant_id, user_id, mock_db):
    """Multiple tool calls in a single LLM response are all executed."""
    handler_a = AsyncMock(return_value="Result A")
    handler_b = AsyncMock(return_value="Result B")
    registry = _make_tool_registry({"tool_a": handler_a, "tool_b": handler_b})

    with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
        mock_llm.side_effect = [
            _make_llm_response(
                content="Calling two tools",
                tool_calls=[
                    _make_tool_call(call_id="c1", name="tool_a", arguments={}),
                    _make_tool_call(call_id="c2", name="tool_b", arguments={}),
                ],
                cost_usd=0.005,
            ),
            _make_llm_response(content="Done with both", cost_usd=0.002),
        ]

        result = await run_react_loop(
            agent_definition=agent_def,
            messages=[{"role": "user", "content": "Call two tools"}],
            tools=[
                {"type": "function", "function": {"name": "tool_a"}},
                {"type": "function", "function": {"name": "tool_b"}},
            ],
            tool_registry=registry,
            db=mock_db,
            tenant_id=tenant_id,
            user_id=user_id,
        )

    assert result.status == "completed"
    assert result.steps_taken == 2
    # Step 1 has two tool calls → two step entries
    assert len(result.steps) == 3  # 2 tool steps + 1 final step
    handler_a.assert_awaited_once()
    handler_b.assert_awaited_once()


@pytest.mark.asyncio
async def test_react_result_dataclass_fields():
    """ReActResult and ReActStep have correct default values."""
    step = ReActStep(
        step_number=1,
        thought="test",
        action=None,
        action_input=None,
        observation=None,
        cost_usd=0.0,
        timestamp="2025-01-01T00:00:00+00:00",
    )
    assert step.step_number == 1
    assert step.thought == "test"

    result = ReActResult(final_content="hello")
    assert result.final_content == "hello"
    assert result.steps == []
    assert result.total_cost_usd == 0.0
    assert result.steps_taken == 0
    assert result.status == "completed"
    assert result.error is None