"""Tests for GraphAugmentedAgent."""

from __future__ import annotations

from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

from koan.agents.types import AgentConfig, AgentResult, RunContext
from koan.graph.augmented import GraphAugmentedAgent
from koan.graph.client import GraphRecord
from koan.graph.context import GraphContext


def _mock_graph_client() -> AsyncMock:
    """Create a mock graph client that returns empty results."""
    client = AsyncMock()
    client.execute = AsyncMock(return_value=[])
    client.execute_write = AsyncMock(return_value=[])
    return client


def _mock_llm_client() -> AsyncMock:
    """Create a mock LLM client."""
    client = AsyncMock()
    response = MagicMock()
    response.text = "I processed your request."
    response.stop_reason = "end_turn"
    response.raw = {}
    response.usage = MagicMock(input_tokens=100, output_tokens=50)
    client.complete = AsyncMock(return_value=response)
    return client


class TestGraphAugmentedAgentRun:
    async def test_basic_run_queries_graph_and_traces(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "test-bot",
            client=llm,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="You are a test agent.",
            ),
            graph_client=graph,
        )

        result = await agent.run("Hello world")

        assert result.output == "I processed your request."
        assert result.agent_name == "test-bot"
        assert "run_id" in result.metadata
        assert "graph_context" in result.metadata

        # Graph should have been written to (tracer writes nodes)
        assert graph.execute_write.call_count > 0

    async def test_run_with_entity_ids(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        # Set up graph to return entity history
        call_count = 0

        async def mock_execute(query: str, params: dict[str, Any] | None = None) -> list[GraphRecord]:
            nonlocal call_count
            call_count += 1
            if "entity_id" in (params or {}):
                return [GraphRecord(data={
                    "entity_id": "acct-001",
                    "entity_type": "account",
                    "decision_count": 3,
                    "last_outcome": "success",
                })]
            return []

        graph.execute = AsyncMock(side_effect=mock_execute)

        agent = GraphAugmentedAgent(
            "finance-bot",
            client=llm,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Handle billing inquiries.",
            ),
            graph_client=graph,
            entity_ids=["acct-001"],
        )

        result = await agent.run("Check billing for acct-001")

        assert result.metadata["graph_context"]["entities"] >= 1
        # Verify the LLM was called with augmented prompt
        llm_call_args = llm.complete.call_args
        messages = llm_call_args[0][1]
        system_msg = messages[0]
        assert "Handle billing inquiries" in system_msg.content

    async def test_run_injects_graph_context_into_prompt(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot",
            client=llm,
            config=AgentConfig(system_prompt="Base prompt."),
            graph_client=graph,
        )

        # Patch gather to return rich context
        context = GraphContext(
            domain_profile={"total_decisions": 42, "decision_types": {"routing": 42}, "outcomes": {"success": 42}},
            active_policies=[{"policy_id": "p1", "name": "Test Policy", "description": "A test"}],
        )
        with patch.object(agent, "_gather_context", return_value=context):
            result = await agent.run("Test input")

        # The LLM should have received the augmented prompt
        llm_call_args = llm.complete.call_args
        messages = llm_call_args[0][1]
        system_content = messages[0].content
        assert "Base prompt." in system_content
        assert "42 past decisions" in system_content
        assert "Test Policy" in system_content

    async def test_run_with_context(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot",
            client=llm,
            config=AgentConfig(),
            graph_client=graph,
        )

        ctx = RunContext(thread_id="t-custom", user_id="alice")
        result = await agent.run("Hello", context=ctx)

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

    async def test_empty_graph_still_works(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot",
            client=llm,
            config=AgentConfig(system_prompt="Just a prompt."),
            graph_client=graph,
        )

        result = await agent.run("Hello")

        # Should work fine with empty graph context (no augmentation)
        assert result.output == "I processed your request."
        llm_call_args = llm.complete.call_args
        messages = llm_call_args[0][1]
        # With empty graph, prompt should just be the base
        assert "Just a prompt." in messages[0].content

    async def test_run_records_entity_links(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot",
            client=llm,
            config=AgentConfig(),
            graph_client=graph,
            entity_ids=["ent-a", "ent-b"],
        )

        await agent.run("Process entities")

        # Check that entity links were written (MERGE Entity + CREATE REFERENCED)
        write_calls = graph.execute_write.call_args_list
        cypher_statements = [call[0][0] for call in write_calls]
        entity_writes = [s for s in cypher_statements if "Entity" in s]
        assert len(entity_writes) >= 2  # One per entity


class TestGraphAugmentedAgentPromptAugmentation:
    def test_augment_prompt_with_context(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot", client=llm, config=AgentConfig(system_prompt="Base."), graph_client=graph,
        )

        ctx = GraphContext(
            domain_profile={"total_decisions": 5, "decision_types": {}, "outcomes": {}},
        )
        result = agent._augment_prompt(ctx)
        assert "Base." in result
        assert "5 past decisions" in result

    def test_augment_prompt_empty_context(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot", client=llm, config=AgentConfig(system_prompt="Base."), graph_client=graph,
        )

        ctx = GraphContext()
        result = agent._augment_prompt(ctx)
        assert result == "Base."

    def test_augment_prompt_no_base_prompt(self) -> None:
        graph = _mock_graph_client()
        llm = _mock_llm_client()

        agent = GraphAugmentedAgent(
            "bot", client=llm, config=AgentConfig(system_prompt=""), graph_client=graph,
        )

        ctx = GraphContext(
            domain_profile={"total_decisions": 1, "decision_types": {}, "outcomes": {}},
        )
        result = agent._augment_prompt(ctx)
        assert "GRAPH CONTEXT" in result
        assert "1 past decisions" in result
