"""Tests for GraphContextProvider and GraphContext."""

from __future__ import annotations

from unittest.mock import AsyncMock

from koan.graph.client import GraphRecord
from koan.graph.context import GraphContext, GraphContextProvider


def _mock_client() -> AsyncMock:
    return AsyncMock()


class TestGraphContext:
    def test_empty_context_returns_empty_string(self) -> None:
        ctx = GraphContext()
        assert ctx.to_prompt() == ""

    def test_domain_profile_only(self) -> None:
        ctx = GraphContext(
            domain_profile={
                "total_decisions": 50,
                "decision_types": {"classification": 30, "routing": 20},
                "outcomes": {"success": 45, "failure": 5},
            }
        )
        prompt = ctx.to_prompt()
        assert "GRAPH CONTEXT" in prompt
        assert "50 past decisions" in prompt
        assert "classification" in prompt
        assert "success" in prompt

    def test_entity_history_formatting(self) -> None:
        ctx = GraphContext(
            entity_history=[
                {
                    "entity_id": "acct-001",
                    "entity_type": "account",
                    "decision_count": 5,
                    "last_outcome": "success",
                },
                {
                    "entity_id": "acct-002",
                    "entity_type": "account",
                    "decision_count": 2,
                    "last_outcome": "failure",
                },
            ]
        )
        prompt = ctx.to_prompt()
        assert "Entity History" in prompt
        assert "acct-001" in prompt
        assert "5 prior decision(s)" in prompt
        assert "acct-002" in prompt

    def test_precedents_formatting(self) -> None:
        ctx = GraphContext(
            precedents=[
                {
                    "description": "Classified as billing",
                    "confidence": 0.95,
                    "outcome_status": "success",
                    "tool_name": "lookup_billing",
                    "reasoning": "Used lookup tool based on account keyword",
                },
            ]
        )
        prompt = ctx.to_prompt()
        assert "Relevant Precedents" in prompt
        assert "Classified as billing" in prompt
        assert "lookup_billing" in prompt
        assert "Used lookup tool" in prompt

    def test_policies_formatting(self) -> None:
        ctx = GraphContext(
            active_policies=[
                {"policy_id": "pol-1", "name": "Data Access Policy", "description": "Controls data access"},
            ]
        )
        prompt = ctx.to_prompt()
        assert "Active Policies" in prompt
        assert "Data Access Policy" in prompt

    def test_full_context(self) -> None:
        ctx = GraphContext(
            domain_profile={"total_decisions": 10, "decision_types": {"routing": 10}, "outcomes": {"success": 10}},
            entity_history=[{"entity_id": "e1", "entity_type": "user", "decision_count": 3, "last_outcome": "success"}],
            precedents=[{"description": "Past decision", "confidence": 0.9, "outcome_status": "success"}],
            active_policies=[{"policy_id": "p1", "name": "Policy A", "description": "Desc"}],
        )
        prompt = ctx.to_prompt()
        assert "Domain Profile" in prompt
        assert "Entity History" in prompt
        assert "Relevant Precedents" in prompt
        assert "Active Policies" in prompt
        assert "END GRAPH CONTEXT" in prompt

    def test_entity_history_capped_at_5(self) -> None:
        ctx = GraphContext(
            entity_history=[
                {"entity_id": f"e-{i}", "entity_type": "t", "decision_count": 1, "last_outcome": "ok"}
                for i in range(10)
            ]
        )
        prompt = ctx.to_prompt()
        assert "e-4" in prompt
        assert "e-5" not in prompt

    def test_precedents_capped_at_5(self) -> None:
        ctx = GraphContext(
            precedents=[
                {"description": f"prec-{i}", "confidence": 0.5, "outcome_status": "ok"}
                for i in range(10)
            ]
        )
        prompt = ctx.to_prompt()
        assert "prec-4" in prompt
        assert "prec-5" not in prompt

    def test_feedback_summary_formatting(self) -> None:
        ctx = GraphContext(
            feedback_summary={
                "total_feedback": 15,
                "positive": 12,
                "negative": 3,
                "satisfaction_rate": 0.8,
                "avg_score": 0.6,
                "recent_comments": ["Great work!", "Could be better"],
            }
        )
        prompt = ctx.to_prompt()
        assert "Feedback History" in prompt
        assert "15 ratings" in prompt
        assert "80%" in prompt
        assert "Great work!" in prompt

    def test_feedback_summary_empty(self) -> None:
        ctx = GraphContext(feedback_summary={})
        assert ctx.to_prompt() == ""


class TestGraphContextProvider:
    async def test_gather_empty(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(return_value=[])

        provider = GraphContextProvider(client)
        ctx = await provider.gather()

        assert ctx.entity_history == []
        assert ctx.precedents == []
        assert ctx.domain_profile == {}
        assert ctx.active_policies == []

    async def test_gather_with_agent_id(self) -> None:
        client = _mock_client()
        # First call: domain profile aggregate
        # Second call: outcome distribution
        # Third call: type distribution
        # Fourth call: active policies
        client.execute = AsyncMock(side_effect=[
            [GraphRecord(data={"total_decisions": 25})],
            [
                GraphRecord(data={"status": "success", "count": 20}),
                GraphRecord(data={"status": "failure", "count": 5}),
            ],
            [
                GraphRecord(data={"dtype": "routing", "count": 15}),
                GraphRecord(data={"dtype": "classification", "count": 10}),
            ],
            [],  # policies
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(agent_id="test-agent")

        assert ctx.domain_profile["total_decisions"] == 25
        assert ctx.domain_profile["outcomes"]["success"] == 20
        assert ctx.domain_profile["decision_types"]["routing"] == 15

    async def test_gather_with_entity_ids(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(side_effect=[
            # entity history for "acct-1"
            [GraphRecord(data={
                "entity_id": "acct-1",
                "entity_type": "account",
                "decision_count": 3,
                "last_outcome": "success",
            })],
            # active policies
            [],
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(entity_ids=["acct-1"])

        assert len(ctx.entity_history) == 1
        assert ctx.entity_history[0]["entity_id"] == "acct-1"

    async def test_gather_with_entity_ids_and_decision_type(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(side_effect=[
            # entity history for "acct-1"
            [GraphRecord(data={
                "entity_id": "acct-1",
                "entity_type": "account",
                "decision_count": 3,
                "last_outcome": "success",
            })],
            # precedents
            [GraphRecord(data={
                "decision_id": "d-1",
                "description": "Past billing classification",
                "reasoning": "Keywords matched",
                "confidence": 0.95,
                "timestamp": "2026-01-01T00:00:00Z",
                "outcome_status": "success",
                "tool_name": "lookup_billing",
            })],
            # active policies
            [],
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(
            entity_ids=["acct-1"],
            decision_type="classification",
        )

        assert len(ctx.precedents) == 1
        assert ctx.precedents[0]["description"] == "Past billing classification"

    async def test_gather_active_policies(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(return_value=[
            GraphRecord(data={
                "policy_id": "pol-1",
                "name": "Data Access",
                "description": "Controls who can access data",
                "version": "1.0",
            }),
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather()

        assert len(ctx.active_policies) == 1
        assert ctx.active_policies[0]["name"] == "Data Access"

    async def test_gather_full(self) -> None:
        client = _mock_client()
        client.execute = AsyncMock(side_effect=[
            # domain profile
            [GraphRecord(data={"total_decisions": 10})],
            # outcomes
            [GraphRecord(data={"status": "success", "count": 10})],
            # type dist
            [GraphRecord(data={"dtype": "routing", "count": 10})],
            # entity history
            [GraphRecord(data={"entity_id": "e1", "entity_type": "user", "decision_count": 2, "last_outcome": "success"})],
            # precedents
            [GraphRecord(data={"decision_id": "d1", "description": "prev", "confidence": 0.8, "outcome_status": "success"})],
            # policies
            [GraphRecord(data={"policy_id": "p1", "name": "P1", "description": "desc"})],
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(
            agent_id="bot",
            entity_ids=["e1"],
            decision_type="routing",
        )

        assert ctx.domain_profile["total_decisions"] == 10
        assert len(ctx.entity_history) == 1
        assert len(ctx.precedents) == 1
        assert len(ctx.active_policies) == 1

        # Full prompt should have all sections
        prompt = ctx.to_prompt()
        assert "Domain Profile" in prompt
        assert "Entity History" in prompt
        assert "Relevant Precedents" in prompt
        assert "Active Policies" in prompt

    async def test_cross_domain_fallback_when_no_entity_precedents(self) -> None:
        """When entity-specific precedents are empty, falls back to cross-domain."""
        client = _mock_client()
        client.execute = AsyncMock(side_effect=[
            # entity history for "acct-new"
            [],
            # entity-specific precedents — empty!
            [],
            # cross-domain precedents fallback
            [GraphRecord(data={
                "decision_id": "d-cross",
                "description": "Cross-domain precedent",
                "confidence": 0.85,
                "outcome_status": "success",
                "tool_name": "lookup_account",
                "entity_ids": ["acct-other"],
            })],
            # active policies
            [],
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(
            entity_ids=["acct-new"],
            decision_type="classification",
        )

        # Should have cross-domain precedents despite new entity
        assert len(ctx.precedents) == 1
        assert ctx.precedents[0]["description"] == "Cross-domain precedent"

    async def test_no_cross_domain_when_entity_precedents_exist(self) -> None:
        """When entity-specific precedents exist, no cross-domain fallback."""
        client = _mock_client()
        client.execute = AsyncMock(side_effect=[
            # entity history
            [GraphRecord(data={
                "entity_id": "acct-1",
                "entity_type": "account",
                "decision_count": 5,
                "last_outcome": "success",
            })],
            # entity-specific precedents — found!
            [GraphRecord(data={
                "decision_id": "d-entity",
                "description": "Entity-specific precedent",
                "confidence": 0.9,
                "outcome_status": "success",
            })],
            # active policies
            [],
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(
            entity_ids=["acct-1"],
            decision_type="classification",
        )

        assert len(ctx.precedents) == 1
        assert ctx.precedents[0]["description"] == "Entity-specific precedent"
        # Only 3 execute calls (no cross-domain fallback)
        assert client.execute.call_count == 3

    async def test_cross_domain_with_decision_type_only(self) -> None:
        """Cross-domain precedents work when no entity_ids provided."""
        client = _mock_client()
        client.execute = AsyncMock(side_effect=[
            # cross-domain precedents
            [GraphRecord(data={
                "decision_id": "d-1",
                "description": "Type-based precedent",
                "confidence": 0.8,
                "outcome_status": "success",
            })],
            # active policies
            [],
        ])

        provider = GraphContextProvider(client)
        ctx = await provider.gather(decision_type="billing_inquiry")

        assert len(ctx.precedents) == 1
        assert ctx.precedents[0]["description"] == "Type-based precedent"
