"""
test_graph_enhanced_retrieval.py
---------------------------------
Test that DualRetriever uses GraphStore for multi-hop expansion.
"""
# nosec B101 - asserts are intentional in test files

import networkx as nx
import numpy as np
from hieromem_core.memory import Memory
from hieromem_core.persistence.in_memory_graph_store import InMemoryGraphStore
from hieromem_core.encoder import Encoder, EncoderConfig


def get_stub_encoder():
    config = EncoderConfig(model_name="stub-hash", backend="stub")
    return Encoder(config)


def test_graph_enhanced_dual_retrieval():
    """
    Test that dual_retrieve leverages GraphStore for multi-hop expansion.

    Scenario:
    - Ep1: Alice knows Bob (vector: [1, 0, 0, 0])
    - Ep2: Bob knows Charlie (vector: [0, 1, 0, 0])
    - Ep3: Charlie knows David (vector: [0, 0, 1, 0])
    - Ep4: Unrelated episode (vector: [0, 0, 0, 1])

    Query: Vector close to Ep1 with query_graph containing "Alice"
    Expected: Should retrieve Ep1 (direct match) and potentially Ep2 (1-hop via graph)
    """
    store = InMemoryGraphStore()
    memory = Memory(store=store, Nmax=10, encoder=get_stub_encoder(), debug=True)

    # Insert Episode 1: Alice knows Bob
    G1 = nx.DiGraph()
    G1.add_edge("Alice", "Bob", relation="knows", trust=0.9)
    vec1 = np.array([1.0, 0.0, 0.0, 0.0])
    ep1 = memory.insert(vec1, G1, meta={"id": "ep1", "desc": "Alice-Bob"})

    # Insert Episode 2: Bob knows Charlie
    G2 = nx.DiGraph()
    G2.add_edge("Bob", "Charlie", relation="knows", trust=0.8)
    vec2 = np.array([0.0, 1.0, 0.0, 0.0])
    ep2 = memory.insert(vec2, G2, meta={"id": "ep2", "desc": "Bob-Charlie"})

    # Insert Episode 3: Charlie knows David
    G3 = nx.DiGraph()
    G3.add_edge("Charlie", "David", relation="knows", trust=0.7)
    vec3 = np.array([0.0, 0.0, 1.0, 0.0])
    memory.insert(vec3, G3, meta={"id": "ep3", "desc": "Charlie-David"})

    # Insert Episode 4: Unrelated
    G4 = nx.DiGraph()
    G4.add_edge("Eve", "Frank", relation="knows", trust=0.6)
    vec4 = np.array([0.0, 0.0, 0.0, 1.0])
    memory.insert(vec4, G4, meta={"id": "ep4", "desc": "Eve-Frank"})

    # Query with vector similar to ep1 and graph containing Alice
    query_vec = np.array([0.95, 0.05, 0.0, 0.0])  # Close to ep1
    query_graph = nx.DiGraph()
    query_graph.add_node("Alice")

    # Perform dual retrieval
    result = memory.dual_retrieve(
        query_vec,
        query_graph=query_graph,
        k=3,
        alpha_dense=1.0,
        beta_symbolic=50.0
    )

    print(f"Retrieved {len(result.episodes)} episodes:")
    for i, ep in enumerate(result.episodes):
        print(f"  {i+1}. {ep.id} - {ep.meta.get('desc')} (weight: {result.weights[i]:.4f})")

    # Assertions
    assert len(result.episodes) > 0, "Should retrieve at least one episode"  # nosec B101

    # Ep1 should be in results (direct vector match + graph match)
    ep_ids = [ep.id for ep in result.episodes]
    assert ep1.id in ep_ids, "Ep1 (Alice-Bob) should be retrieved"  # nosec B101

    # With graph expansion, Ep2 might also be retrieved (1-hop from Alice via Bob)
    # This is not guaranteed but demonstrates the capability
    if ep2.id in ep_ids:
        print("✓ Graph expansion successfully retrieved Ep2 (Bob-Charlie) via 1-hop from Alice")

    print("✓ Graph-enhanced dual retrieval test passed!")


if __name__ == "__main__":
    test_graph_enhanced_dual_retrieval()
