"""Tests for multi-agent feedback propagation."""

from __future__ import annotations

from unittest.mock import AsyncMock

from koan.graph.client import GraphRecord
from koan.graph.propagation import FeedbackPropagator, PropagationResult, link_runs


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


class TestLinkRuns:
    async def test_link_creates_edge(self) -> None:
        client = _mock_client()
        await link_runs(client, parent_run_id="p1", child_run_id="c1")
        assert client.execute_write.call_count == 1
        query = client.execute_write.call_args[0][0]
        assert "DELEGATED_TO" in query
        params = client.execute_write.call_args[0][1]
        assert params["parent_id"] == "p1"
        assert params["child_id"] == "c1"

    async def test_link_custom_relationship(self) -> None:
        client = _mock_client()
        await link_runs(client, parent_run_id="p1", child_run_id="c1", relationship="HANDED_OFF")
        query = client.execute_write.call_args[0][0]
        assert "HANDED_OFF" in query


class TestFeedbackPropagator:
    async def test_no_upstream_runs(self) -> None:
        client = _mock_client()
        prop = FeedbackPropagator(client)
        result = await prop.propagate(run_id="r1", score=1.0)
        assert result.source_run_id == "r1"
        assert result.propagated_to == []
        assert result.feedback_ids == []

    async def test_single_upstream_run(self) -> None:
        client = _mock_client()
        # First call: find upstream runs
        # Second call: record feedback
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={"run_id": "parent-run"}),
        ])
        prop = FeedbackPropagator(client)
        result = await prop.propagate(run_id="child-run", score=1.0, comment="Great!")
        assert len(result.propagated_to) == 1
        assert result.propagated_to[0] == "parent-run"
        assert len(result.feedback_ids) == 1
        # Score should be decayed
        write_call = client.execute_write.call_args
        assert write_call[0][1]["score"] == 0.7  # 1.0 * 0.7

    async def test_multi_hop_decay(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={"run_id": "parent-run"}),
            GraphRecord(data={"run_id": "grandparent-run"}),
        ])
        prop = FeedbackPropagator(client, decay_factor=0.5)
        result = await prop.propagate(run_id="child", score=1.0)
        assert len(result.propagated_to) == 2

        # First write: parent gets 0.5
        first_write = client.execute_write.call_args_list[0]
        assert first_write[0][1]["score"] == 0.5

        # Second write: grandparent gets 0.25
        second_write = client.execute_write.call_args_list[1]
        assert second_write[0][1]["score"] == 0.25

    async def test_negative_score_propagates(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={"run_id": "parent"}),
        ])
        prop = FeedbackPropagator(client)
        result = await prop.propagate(run_id="child", score=-1.0)
        write_params = client.execute_write.call_args[0][1]
        assert write_params["score"] == -0.7

    async def test_max_hops_respected(self) -> None:
        client = _mock_client()
        # Return 10 upstream runs
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={"run_id": f"run-{i}"})
            for i in range(10)
        ])
        prop = FeedbackPropagator(client, max_hops=3)
        result = await prop.propagate(run_id="child", score=1.0)
        assert len(result.propagated_to) == 3

    async def test_propagated_feedback_is_marked(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={"run_id": "parent"}),
        ])
        prop = FeedbackPropagator(client)
        await prop.propagate(run_id="child", score=1.0, comment="Nice")
        write_params = client.execute_write.call_args[0][1]
        assert write_params["feedback_type"] == "propagated"
        assert "Propagated" in write_params["comment"]

    async def test_get_chain_empty(self) -> None:
        client = _mock_client()
        prop = FeedbackPropagator(client)
        chain = await prop.get_chain("run-1")
        assert chain == []

    async def test_get_chain_with_data(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={"run_id": "r1", "agent_id": "agent-a", "status": "completed"}),
            GraphRecord(data={"run_id": "r2", "agent_id": "agent-b", "status": "completed"}),
        ])
        prop = FeedbackPropagator(client)
        chain = await prop.get_chain("r2")
        assert len(chain) == 2
        assert chain[0]["agent_id"] == "agent-a"


class TestPropagationResult:
    def test_defaults(self) -> None:
        r = PropagationResult(source_run_id="r1")
        assert r.propagated_to == []
        assert r.feedback_ids == []
