"""Tests for traced_agent() middleware wrapping KOAN agents."""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

from koan.agents.types import AgentResult, RunContext
from koan.graph.client import GraphClient
from koan.graph.middleware import traced_agent


def _mock_graph_client() -> GraphClient:
    client = AsyncMock(spec=GraphClient)
    client.execute_write = AsyncMock(return_value=[])
    client.execute = AsyncMock(return_value=[])
    return client


def _mock_agent(
    name: str = "test-agent",
    output: str = "Hello!",
    tool_calls: int = 0,
    error: bool = False,
) -> MagicMock:
    """Create a mock agent with a preset result."""
    result = AgentResult(
        output=output,
        agent_name=name,
        iterations=1,
        tool_calls_made=tool_calls,
        metadata={"error": "some error"} if error else {},
    )
    agent = MagicMock()
    agent.name = name
    agent.run = AsyncMock(return_value=result)
    return agent


class TestTracedAgent:
    async def test_runs_agent_and_returns_result(self) -> None:
        agent = _mock_agent(output="Test output")
        graph = _mock_graph_client()

        result = await traced_agent(agent, "Hello", client=graph)

        assert result.output == "Test output"
        agent.run.assert_awaited_once()

    async def test_passes_context_to_agent(self) -> None:
        agent = _mock_agent()
        graph = _mock_graph_client()
        ctx = RunContext(user_id="u-1", thread_id="t-1")

        await traced_agent(agent, "Hello", client=graph, context=ctx)

        agent.run.assert_awaited_once_with("Hello", context=ctx)

    async def test_creates_run_and_outcome_in_graph(self) -> None:
        agent = _mock_agent()
        graph = _mock_graph_client()

        await traced_agent(agent, "Hello", client=graph)

        # Should have at minimum: __aenter__(Run), decision, outcome, __aexit__
        assert graph.execute_write.call_count >= 4

        # First call creates Run node
        enter_cypher = graph.execute_write.call_args_list[0][0][0]
        assert "CREATE (r:Run" in enter_cypher

        # Last call is __aexit__ setting status
        exit_cypher = graph.execute_write.call_args_list[-1][0][0]
        assert "SET r.status" in exit_cypher

    async def test_records_routing_decision(self) -> None:
        agent = _mock_agent(name="my-agent")
        graph = _mock_graph_client()

        await traced_agent(agent, "Hello", client=graph)

        decision_call = graph.execute_write.call_args_list[1]
        cypher = decision_call[0][0]
        params = decision_call[0][1]
        assert "CREATE (d:Decision" in cypher
        assert params["decision_type"] == "routing"
        assert "my-agent" in params["description"]

    async def test_records_tool_action_when_tools_used(self) -> None:
        agent = _mock_agent(tool_calls=3)
        graph = _mock_graph_client()

        await traced_agent(agent, "Hello", client=graph)

        # Should have: __aenter__, decision, tool action, outcome, __aexit__
        assert graph.execute_write.call_count == 5
        action_call = graph.execute_write.call_args_list[2]
        cypher = action_call[0][0]
        params = action_call[0][1]
        assert "CREATE (a:Action" in cypher
        assert params["type"] == "tool_call"

    async def test_no_tool_action_when_no_tools(self) -> None:
        agent = _mock_agent(tool_calls=0)
        graph = _mock_graph_client()

        await traced_agent(agent, "Hello", client=graph)

        # Should have: __aenter__, decision, outcome, __aexit__ (no action)
        assert graph.execute_write.call_count == 4

    async def test_records_failure_outcome(self) -> None:
        agent = _mock_agent(error=True)
        graph = _mock_graph_client()

        await traced_agent(agent, "Hello", client=graph)

        # Outcome call — second-to-last (before __aexit__)
        outcome_call = graph.execute_write.call_args_list[-2]
        params = outcome_call[0][1]
        assert params["status"] == "failure"

    async def test_records_success_outcome(self) -> None:
        agent = _mock_agent(error=False)
        graph = _mock_graph_client()

        await traced_agent(agent, "Hello", client=graph)

        outcome_call = graph.execute_write.call_args_list[-2]
        params = outcome_call[0][1]
        assert params["status"] == "success"

    async def test_attaches_run_id_to_result_metadata(self) -> None:
        agent = _mock_agent()
        graph = _mock_graph_client()

        result = await traced_agent(agent, "Hello", client=graph)

        assert "run_id" in result.metadata
        assert "thread_id" in result.metadata

    async def test_uses_context_thread_id(self) -> None:
        agent = _mock_agent()
        graph = _mock_graph_client()
        ctx = RunContext(thread_id="my-thread")

        result = await traced_agent(agent, "Hello", client=graph, context=ctx)

        assert result.metadata["thread_id"] == "my-thread"

    async def test_uses_explicit_thread_id(self) -> None:
        agent = _mock_agent()
        graph = _mock_graph_client()

        result = await traced_agent(
            agent, "Hello", client=graph, thread_id="explicit-t"
        )

        assert result.metadata["thread_id"] == "explicit-t"

    async def test_explicit_thread_id_overrides_context(self) -> None:
        agent = _mock_agent()
        graph = _mock_graph_client()
        ctx = RunContext(thread_id="from-context")

        result = await traced_agent(
            agent, "Hello", client=graph, context=ctx, thread_id="explicit"
        )

        assert result.metadata["thread_id"] == "explicit"
