"""
Tests for the retrieve_context LangGraph node.

These tests focus on:
- Combining knowledge-base and user-memory context
- Retrieving user memory only on the first turn
- Normalizing both sources into the same retrieved_context shape

To run this test, enter the following code in the terminal:
    PYTHONPATH=src pytest tests/test_retrieve_context_node.py -q -s
"""

from __future__ import annotations
from walkxr_ai.core.nodes import build_retrieve_context_node


class FakeRetrievalEngine:
    """
    Fake knowledgebase retriever returning a fixed normalized result.
    """

    def retrieve(self, query: str) -> list[dict]:
        return [
            {
                "text": "Knowledge: grounding techniques help with anxiety.",
                "source": "knowledge_base",
                "score": 0.9,
                "metadata": {},
            }
        ]


class FakeNode:
    """
    Fake node that mimics the LlamaIndex interface used by the node.
    """

    def __init__(self, text: str, metadata: dict) -> None:
        self._text = text
        self.metadata = metadata

    def get_text(self) -> str:
        return self._text


class FakeResult:
    """
    Fake result that mimics a LlamaIndex retrieval result.
    """

    def __init__(self, text: str) -> None:
        self.node = FakeNode(text, {"user_id": "user_1"})
        self.score = 0.8


class FakeUserMemoryStore:
    """
    Fake user-memory store returning LlamaIndex-like result objects.
    """

    def query_memories(self, user_id: str, query: str):
        return [FakeResult("Memory: the user finds reflective walks relaxing.")]


def test_retrieve_context_combines_knowledge_and_memory_on_new_session() -> None:
    """
    Test #1:
    Test to verify that the knowledgebase (non-memory) context and user memory context 
    are combined on the first turn.
    """
    retrieve_context = build_retrieve_context_node(
        FakeRetrievalEngine(),
        FakeUserMemoryStore(),
    )

    state = {
        "user_input": "How can I relax?",
        "user_id": "user_1",
        "turn_index": 0,
    }

    result = retrieve_context(state)

    assert result["phase"] == "context_retrieved"                                           # Is the phase correct?
    assert any(item["source"] == "knowledge_base" for item in result["retrieved_context"])  # Is knowledgebase context included?
    assert any(item["source"] == "user_memory" for item in result["retrieved_context"])     # Is user memory context included?


def test_retrieve_context_skips_memory_after_first_turn() -> None:
    """
    Test #2:
    Test to verify that user-memory retrieval is skipped on subsequent turns after the first 
    turn of the session, even if user_id is present. This ensures that we only retrieve user 
    memory on the first turn, as intended.
    """
    retrieve_context = build_retrieve_context_node(
        FakeRetrievalEngine(),
        FakeUserMemoryStore(),
    )

    state = {
        "user_input": "How can I relax?",
        "user_id": "user_1",
        "turn_index": 1,
    }

    result = retrieve_context(state)

    assert result["phase"] == "context_retrieved"                                           # Is the phase correct?
    assert len(result["retrieved_context"]) == 1                                            # Is only one context item returned (the knowledgebase context)?
    assert all(item["source"] != "user_memory" for item in result["retrieved_context"])     # Is user memory context skipped?


def test_retrieve_context_skips_memory_if_no_user_id() -> None:
    """
    Test #3:
    Test to verify that user-memory retrieval is skipped if user_id is blank, even on the first turn. 
    This ensures that we only retrieve user memory when we have a valid user_id, as intended.
    """
    retrieve_context = build_retrieve_context_node(
        FakeRetrievalEngine(),
        FakeUserMemoryStore(),
    )

    state = {
        "user_input": "How can I relax?",
        "user_id": "",
        "turn_index": 0,
    }

    result = retrieve_context(state)

    assert result["phase"] == "context_retrieved"                                           # Is the phase correct?
    assert len(result["retrieved_context"]) == 1                                            # Is only one context item returned (the knowledgebase context)?
    assert all(item["source"] != "user_memory" for item in result["retrieved_context"])     # Is user memory context skipped?


def test_retrieve_context_returns_expected_structure() -> None:
    """
    Test #4:
    Test to verify that the retrieved context items from both knowledgebase and user memory 
    are returned in the expected normalized structure, with keys: text, source, score, and metadata. 
    This ensures that the node is correctly normalizing different context sources into a consistent 
    format for downstream nodes.
    """
    retrieve_context = build_retrieve_context_node(
        FakeRetrievalEngine(),
        FakeUserMemoryStore(),
    )

    state = {
        "user_input": "How can I relax?",
        "user_id": "user_1",
        "turn_index": 0,
    }

    result = retrieve_context(state)

    assert result["phase"] == "context_retrieved"   # Is the phase correct?
    for item in result["retrieved_context"]:        # Does each item have the expected keys?
        assert "text" in item
        assert "source" in item
        assert "score" in item
        assert "metadata" in item