"""Tests for DecisionTracer against a mocked GraphClient."""

from __future__ import annotations

from unittest.mock import AsyncMock

from koan.graph.client import GraphClient
from koan.graph.tracer import DecisionTracer


def _mock_graph_client() -> GraphClient:
    """Create a mock GraphClient that records execute_write calls."""
    client = AsyncMock(spec=GraphClient)
    client.execute_write = AsyncMock(return_value=[])
    client.execute = AsyncMock(return_value=[])
    return client


class TestDecisionTracerEnterExit:
    async def test_aenter_creates_run_node(self) -> None:
        client = _mock_graph_client()
        tracer = DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1")

        async with tracer:
            pass

        # First call is __aenter__ creating the Run node
        enter_call = client.execute_write.call_args_list[0]
        cypher = enter_call[0][0]
        params = enter_call[0][1]
        assert "CREATE (r:Run" in cypher
        assert "MERGE (a:Agent" in cypher
        assert params["run_id"] == "r-1"
        assert params["agent_id"] == "a-1"
        assert params["thread_id"] == "t-1"

    async def test_aexit_sets_completed_status(self) -> None:
        client = _mock_graph_client()
        tracer = DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1")

        async with tracer:
            pass

        # Last call is __aexit__ updating status
        exit_call = client.execute_write.call_args_list[-1]
        cypher = exit_call[0][0]
        params = exit_call[0][1]
        assert "SET r.status" in cypher
        assert params["status"] == "completed"

    async def test_aexit_sets_failed_on_exception(self) -> None:
        client = _mock_graph_client()
        tracer = DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1")

        try:
            async with tracer:
                msg = "boom"
                raise RuntimeError(msg)
        except RuntimeError:
            pass

        exit_call = client.execute_write.call_args_list[-1]
        params = exit_call[0][1]
        assert params["status"] == "failed"


class TestRecordDecision:
    async def test_creates_decision_linked_to_run(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            decision_id = await tracer.record_decision(
                description="Classified as billing",
                decision_type="classification",
                confidence=0.92,
                reasoning="Keyword match",
            )

        assert decision_id  # non-empty UUID string

        # Find the decision creation call (second call, after __aenter__)
        decision_call = client.execute_write.call_args_list[1]
        cypher = decision_call[0][0]
        params = decision_call[0][1]
        assert "CREATE (d:Decision" in cypher
        assert "(r)-[:CONTAINS]->(d)" in cypher
        assert params["description"] == "Classified as billing"
        assert params["decision_type"] == "classification"
        assert params["confidence"] == 0.92

    async def test_links_to_policy_when_provided(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            await tracer.record_decision(
                description="Evaluated refund policy",
                decision_type="policy_eval",
                policy_id="pol-1",
            )

        # Should have 4 calls: __aenter__, decision, policy link, __aexit__
        assert client.execute_write.call_count == 4
        policy_call = client.execute_write.call_args_list[2]
        cypher = policy_call[0][0]
        params = policy_call[0][1]
        assert "GOVERNED_BY" in cypher
        assert params["policy_id"] == "pol-1"

    async def test_links_to_precedent_when_provided(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            await tracer.record_decision(
                description="Cited previous decision",
                decision_type="exception",
                precedent_decision_id="prev-d-1",
            )

        # Should have 4 calls: __aenter__, decision, precedent link, __aexit__
        assert client.execute_write.call_count == 4
        precedent_call = client.execute_write.call_args_list[2]
        cypher = precedent_call[0][0]
        params = precedent_call[0][1]
        assert "CITED_PRECEDENT" in cypher
        assert params["prev_id"] == "prev-d-1"


class TestRecordAction:
    async def test_creates_action_linked_to_decision(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            decision_id = await tracer.record_decision(
                description="Selected tool",
                decision_type="tool_selection",
            )
            action_id = await tracer.record_action(
                decision_id=decision_id,
                action_type="tool_call",
                tool_name="web_search",
                latency_ms=340,
            )

        assert action_id
        action_call = client.execute_write.call_args_list[2]
        cypher = action_call[0][0]
        params = action_call[0][1]
        assert "CREATE (a:Action" in cypher
        assert "(d)-[:LED_TO]->(a)" in cypher
        assert params["type"] == "tool_call"
        assert params["tool_name"] == "web_search"
        assert params["latency_ms"] == 340

    async def test_records_error_in_action(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            decision_id = await tracer.record_decision(
                description="Tool call",
                decision_type="tool_selection",
            )
            await tracer.record_action(
                decision_id=decision_id,
                action_type="tool_call",
                tool_name="failing_tool",
                error="Connection timeout",
            )

        action_call = client.execute_write.call_args_list[2]
        params = action_call[0][1]
        assert params["error"] == "Connection timeout"


class TestRecordOutcome:
    async def test_creates_outcome_linked_to_run(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            outcome_id = await tracer.record_outcome(
                status="success",
                result="Refund processed",
                evaluation_score=0.95,
            )

        assert outcome_id
        outcome_call = client.execute_write.call_args_list[1]
        cypher = outcome_call[0][0]
        params = outcome_call[0][1]
        assert "CREATE (o:Outcome" in cypher
        assert "(r)-[:PRODUCED]->(o)" in cypher
        assert params["status"] == "success"
        assert params["eval_score"] == 0.95


class TestLinkEntity:
    async def test_links_entity_to_decision(self) -> None:
        client = _mock_graph_client()

        async with DecisionTracer(client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            decision_id = await tracer.record_decision(
                description="Reviewed account",
                decision_type="classification",
            )
            await tracer.link_entity(
                decision_id=decision_id,
                entity_id="acct-123",
                entity_type="account",
                source_system="salesforce",
            )

        entity_call = client.execute_write.call_args_list[2]
        cypher = entity_call[0][0]
        params = entity_call[0][1]
        assert "MERGE (e:Entity" in cypher
        assert "(d)-[:REFERENCED]->(e)" in cypher
        assert params["entity_id"] == "acct-123"
        assert params["entity_type"] == "account"
        assert params["source"] == "salesforce"
