"""
test_graph_sanity.py
---------------------
Phase 3.1: Graph sanity integration test.
Verifies the graph layer provides real signal on actual data without Neo4j.

Phase 3.2: Dual retrieve recall check.
Measures if dual retrieval improves recall over vector-only retrieval.
"""
# 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():
    """Create a stub encoder for testing."""
    config = EncoderConfig(model_name="stub-hash", backend="stub")
    return Encoder(config)


def extract_entities(text: str) -> list[str]:
    """Simple deterministic entity extraction: capitalized words."""
    words = text.split()
    entities = []
    for word in words:
        clean = word.strip(".,!?;:'\"")
        if clean and clean[0].isupper() and len(clean) > 1:
            entities.append(clean)
    return list(set(entities))


def create_graph_from_text(text: str, episode_id: str) -> nx.DiGraph:
    """Create a graph with co-occurrence edges between entities."""
    entities = extract_entities(text)
    G = nx.DiGraph()

    # Add mentions edges
    for entity in entities:
        G.add_edge(episode_id, entity, relation="mentions")

    # Add co-occurrence edges between entities
    for i, e1 in enumerate(entities):
        for e2 in entities[i + 1:]:
            G.add_edge(e1, e2, relation="co_occurs_with")

    return G


# Sample paragraphs for testing (simplified HotPot-style)
SAMPLE_PARAGRAPHS = [
    "Albert Einstein was born in Ulm, Germany. He later moved to Switzerland.",
    "Switzerland is known for its mountains and the city of Bern.",
    "Bern is the capital of Switzerland and has beautiful architecture.",
    "Germany is a European country known for its engineering.",
    "Ulm is a city in Germany on the Danube River.",
    "The Danube River flows through many European countries.",
    "Physics was revolutionized by Albert Einstein's theories.",
    "Relativity theory was developed by Albert Einstein in 1905.",
    "Isaac Newton developed classical mechanics before Einstein.",
    "Newton was born in England and studied at Cambridge.",
    "Cambridge University is one of the oldest universities.",
    "England is part of the United Kingdom.",
    "Marie Curie won Nobel Prizes in Physics and Chemistry.",
    "Chemistry and Physics are fundamental sciences.",
    "Paris is where Marie Curie conducted her research.",
    "France is home to Paris and many cultural sites.",
]


def test_graph_sanity_integration():
    """
    Phase 3.1: Graph sanity integration test.

    1. Insert paragraphs with entity-based graphs
    2. Run graph_search from a seed
    3. Verify we get related episodes through entity connections
    """
    store = InMemoryGraphStore()
    memory = Memory(store=store, Nmax=50, encoder=get_stub_encoder(), debug=False)

    episode_ids = []

    # Insert all paragraphs with graph data
    for i, para in enumerate(SAMPLE_PARAGRAPHS):
        # Create embedding based on hash of text (deterministic)
        vec = np.array([hash(para) % 1000 / 1000.0 for _ in range(4)])
        G = create_graph_from_text(para, f"ep_{i}")

        ep = memory.insert(vec, G, meta={"id": f"ep_{i}", "text": para})
        episode_ids.append(ep.id)

    # Verify graph was populated
    stats = memory.stats()
    assert stats["graph_store"]["enabled"] is True, "Graph should be enabled"  # nosec B101
    assert stats["graph_store"]["node_count"] > 0, "Should have nodes"  # nosec B101
    assert stats["graph_store"]["edge_count"] > 0, "Should have edges"  # nosec B101

    print(f"Graph stats: {stats['graph_store']['node_count']} nodes, {stats['graph_store']['edge_count']} edges")

    # Pick seed episode (Einstein paragraph)
    seed_id = episode_ids[0]  # "Albert Einstein was born in Ulm, Germany..."

    # Run graph search with depth=2
    results = memory.graph_search(seed_id, depth=2, limit=50)

    # Verify we found related content
    assert len(results) > 0, "Should find related entities"  # nosec B101

    # Check if we found connections to other episodes through shared entities
    found_entities = set()
    for r in results:
        found_entities.add(r.get("source"))
        found_entities.add(r.get("target"))

    # We should find entities that connect to other episodes
    # e.g., "Germany" should connect Einstein paragraph to other Germany-related content
    print(f"Found {len(found_entities)} unique entities through graph traversal")
    assert "Germany" in found_entities or "Einstein" in found_entities, \
        "Should find key entities from seed"  # nosec B101

    # Test graph_diagnostics
    diag = memory.graph_diagnostics()
    assert diag["enabled"] is True  # nosec B101
    assert "top_relations" in diag  # nosec B101
    print(f"Top relations: {diag['top_relations']}")

    print("✓ Phase 3.1: Graph sanity integration test PASSED")


def test_dual_retrieve_improves_recall():
    """
    Phase 3.2: Dual retrieve recall check.

    Measures if dual retrieval finds more relevant episodes than vector-only.
    """
    store = InMemoryGraphStore()
    memory = Memory(store=store, Nmax=50, encoder=get_stub_encoder(), debug=False)

    episode_ids = []
    episode_entities: dict[str, set[str]] = {}

    # Insert paragraphs with graph data
    for i, para in enumerate(SAMPLE_PARAGRAPHS):
        vec = np.array([hash(para) % 1000 / 1000.0 for _ in range(4)])
        G = create_graph_from_text(para, f"ep_{i}")

        ep = memory.insert(vec, G, meta={"id": f"ep_{i}", "text": para})
        episode_ids.append(ep.id)
        episode_entities[ep.id] = set(extract_entities(para))

    # Query: Looking for Einstein-related content
    query_text = "Albert Einstein physics relativity"
    query_entities = extract_entities(query_text)
    query_vec = np.array([hash(query_text) % 1000 / 1000.0 for _ in range(4)])

    # Define gold set: episodes sharing entities with query
    gold_ids = set()
    for ep_id, entities in episode_entities.items():
        if entities & set(query_entities):
            gold_ids.add(ep_id)

    print(f"Gold episode IDs (share entities with query): {gold_ids}")

    # Vector-only retrieval
    vector_results = memory.resonant_retrieve(query_vec, k=5)
    vector_ids = {ep.id for ep in vector_results}

    # Dual retrieval with graph
    query_graph = nx.DiGraph()
    for entity in query_entities:
        query_graph.add_node(entity)

    dual_result = memory.dual_retrieve(
        query_vec,
        query_graph=query_graph,
        k=5,
        alpha_dense=1.0,
        beta_symbolic=50.0
    )
    dual_ids = {ep.id for ep in dual_result.episodes}

    # Calculate recall
    vector_recall = len(vector_ids & gold_ids) / max(len(gold_ids), 1)
    dual_recall = len(dual_ids & gold_ids) / max(len(gold_ids), 1)

    print(f"Vector-only recall: {vector_recall:.2%}")
    print(f"Dual retrieval recall: {dual_recall:.2%}")

    # Dual should be at least as good as vector-only
    # (The stub encoder produces hash-based embeddings, so graph may not always help)
    # We allow equality but not significant degradation
    assert dual_recall >= vector_recall, \
        f"Dual recall ({dual_recall:.2%}) should be >= vector recall ({vector_recall:.2%})"  # nosec B101

    print("✓ Phase 3.2: Dual retrieve recall check PASSED")


def test_summarise_preserves_graph():
    """
    Test that summarisation properly handles the graph layer.

    When episodes are merged during summarisation, the graph store
    should retain the entities from the merged episodes.
    """
    store = InMemoryGraphStore()
    # Set Nmax low to trigger summarisation
    memory = Memory(store=store, Nmax=3, encoder=get_stub_encoder(), debug=False)

    # Create 4 episodes - will trigger summarisation after 3rd
    episodes = []
    for i in range(4):
        text = f"Entity{i} is related to Common and also to Shared{i}."
        G = create_graph_from_text(text, f"ep_{i}")
        vec = np.random.rand(128).astype(np.float32)
        ep = memory.insert(vec, G, meta={"text": text})
        episodes.append(ep)

    # Check that graph store still has entries
    stats = store.get_graph_stats()
    assert stats["enabled"] is True  # nosec B101
    assert stats["node_count"] > 0, "Graph should have nodes after summarisation"  # nosec B101
    assert stats["edge_count"] > 0, "Graph should have edges after summarisation"  # nosec B101

    # The graph should preserve entities from ALL episodes (original + summary)
    # Query for a common entity
    relations = store.query_relations("Common", depth=1)
    # Should find connections since all episodes mentioned "Common"
    print(f"Found {len(relations)} relations from 'Common' after summarisation")

    print("✓ test_summarise_preserves_graph PASSED")


def test_concurrent_ingestion_warning():
    """
    Test that documents the thread-safety limitation.

    This test verifies the InMemoryGraphStore docstring warns about
    thread safety, since NetworkX is not thread-safe.
    """
    from hieromem_core.persistence.in_memory_graph_store import InMemoryGraphStore

    # Check that the docstring contains the warning
    docstring = InMemoryGraphStore.__doc__ or ""
    assert "NOT thread-safe" in docstring, \
        "InMemoryGraphStore should document thread-safety warning"  # nosec B101
    assert "NetworkX" in docstring, \
        "Warning should mention NetworkX limitation"  # nosec B101
    assert "Neo4jStore" in docstring, \
        "Warning should recommend Neo4jStore for concurrent access"  # nosec B101

    print("✓ test_concurrent_ingestion_warning PASSED")


def test_concurrent_read_write_stability():
    """
    Test concurrent read/write stability with threading.

    Start one thread inserting episodes with graphs repeatedly.
    Start another calling search_related() in a loop.
    Assert: no exceptions + edge counts increase monotonically.
    """
    import threading
    import time

    store = InMemoryGraphStore()
    memory = Memory(store=store, Nmax=100, encoder=get_stub_encoder(), debug=False)

    errors: list[Exception] = []
    edge_counts: list[int] = []

    def writer_thread():
        """Insert episodes with graphs repeatedly."""
        for i in range(10):
            try:
                G = nx.DiGraph()
                G.add_edge(f"Entity_{i}", "SharedNode", relation="connects")
                vec = np.random.rand(128).astype(np.float32)
                memory.insert(vec, G, meta={"thread": "writer", "iteration": i})
                time.sleep(0.01)  # Small delay
            except Exception as e:
                errors.append(e)

    def reader_thread():
        """Query search_related() in a loop."""
        for i in range(20):
            try:
                # Try to search - may return empty if no episodes yet
                if memory.episodes:
                    ep = memory.episodes[0]
                    store.search_related(ep.id, depth=2, limit=10)
                stats = store.get_graph_stats()
                edge_counts.append(stats["edge_count"])
                time.sleep(0.005)  # Faster reads
            except Exception as e:
                errors.append(e)

    # Start threads
    writer = threading.Thread(target=writer_thread)
    reader = threading.Thread(target=reader_thread)

    writer.start()
    reader.start()

    writer.join()
    reader.join()

    # Assert no exceptions
    assert len(errors) == 0, f"Concurrent access raised exceptions: {errors}"  # nosec B101

    # Edge counts should be monotonically non-decreasing (within reader snapshots)
    # Some slight reordering is possible, but no major drops
    if len(edge_counts) > 2:
        # Allow small fluctuations but no major drops
        max_count = max(edge_counts)
        final_count = edge_counts[-1]
        assert final_count >= max_count * 0.5, \
            f"Edge count dropped unexpectedly: {edge_counts}"  # nosec B101

    print(f"Concurrent test: {len(edge_counts)} reads, final edge count: {edge_counts[-1] if edge_counts else 0}")
    print("✓ test_concurrent_read_write_stability PASSED")


def test_summarise_no_duplicate_edges():
    """
    Test that summarisation does not create duplicate edges.

    Create 2 episodes with identical edge A→B.
    Force summarise.
    Assert: summary contributes expected number of edges once.
    """
    store = InMemoryGraphStore()
    # Set Nmax=2 so adding 3 episodes triggers summarisation
    memory = Memory(store=store, Nmax=2, encoder=get_stub_encoder(), debug=False)

    # Create 2 episodes with identical edges
    G1 = nx.DiGraph()
    G1.add_edge("A", "B", relation="identical")
    vec1 = np.random.rand(128).astype(np.float32)
    memory.insert(vec1, G1, meta={"ep": 1})

    G2 = nx.DiGraph()
    G2.add_edge("A", "B", relation="identical")
    vec2 = np.random.rand(128).astype(np.float32)
    memory.insert(vec2, G2, meta={"ep": 2})

    # Get edge count before summarisation trigger
    stats_before = store.get_graph_stats()
    edge_count_before = stats_before["edge_count"]
    print(f"Before 3rd insert: {edge_count_before} edges")

    # Insert 3rd episode to trigger summarisation
    G3 = nx.DiGraph()
    G3.add_edge("C", "D", relation="different")
    vec3 = np.random.rand(128).astype(np.float32)
    memory.insert(vec3, G3, meta={"ep": 3})

    # Get edge count after summarisation
    stats_after = store.get_graph_stats()
    edge_count_after = stats_after["edge_count"]
    print(f"After summarisation: {edge_count_after} edges")

    # Edge count should not blow up
    # Expected: original edges (2 from ep1, ep2) + new edges from ep3 and summary
    # But NOT doubled due to double-save
    # Maximum reasonable: 2 (original) + 1 (new) + some summary edges
    assert edge_count_after < edge_count_before * 4, \
        f"Edge count blew up after summarisation: {edge_count_before} -> {edge_count_after}"  # nosec B101

    # More specifically: edge count should not have doubled unexpectedly
    # If double-save bug exists, we'd see 2x edges
    print(f"Edge count ratio: {edge_count_after / max(edge_count_before, 1):.2f}x")

    print("✓ test_summarise_no_duplicate_edges PASSED")


def test_graph_is_persisted():
    """
    Phase 1 Invariant: After insert, graph MUST be in graph_store.

    This test verifies that when you insert an episode with a graph,
    the graph is persisted to graph_store and is queryable.
    """
    store = InMemoryGraphStore()
    memory = Memory(store=store, Nmax=100, encoder=get_stub_encoder(), debug=False)

    # Verify graph_store is NEVER None
    assert memory.graph_store is not None, "graph_store must NEVER be None"  # nosec B101

    # Insert episode with graph
    G = nx.DiGraph()
    G.add_edge("Alice", "Bob", relation="knows")
    G.add_edge("Bob", "Charlie", relation="works_with")

    vec = np.random.rand(384).astype(np.float32)
    memory.insert(vec, G, meta={"title": "test"}, namespace="test_ns")

    # Graph MUST be persisted
    stats = memory.graph_store.get_graph_stats(namespace="test_ns")
    assert stats["edge_count"] >= 2, f"Expected at least 2 edges, got {stats['edge_count']}"  # nosec B101
    assert stats["node_count"] >= 3, f"Expected at least 3 nodes, got {stats['node_count']}"  # nosec B101

    print("✓ test_graph_is_persisted PASSED")


def test_gsr_filters_not_erases():
    """
    Phase 3 Invariant: GSR filters edges, never prevents memory formation.

    This test verifies that when strict_gsr=False, forbidden edges are
    dropped but the episode still persists with allowed edges.
    """
    store = InMemoryGraphStore()
    memory = Memory(
        store=store, Nmax=100, encoder=get_stub_encoder(),
        debug=False, strict_gsr=False  # KEY: filter, don't raise
    )

    # Graph with mixed edges: some forbidden, some allowed
    G = nx.DiGraph()
    G.add_edge("Alice", "Bob", relation="knows")           # Allowed
    G.add_edge("user", "data", relation="has_secret")      # Forbidden
    G.add_edge("Bob", "Charlie", relation="works_with")    # Allowed

    vec = np.random.rand(384).astype(np.float32)
    memory.insert(vec, G, meta={"title": "test"}, namespace="test_ns")

    # Episode MUST be persisted
    assert len(memory.episodes) == 1, "Episode must be persisted"  # nosec B101

    # Graph should have allowed edges only
    stats = memory.graph_store.get_graph_stats(namespace="test_ns")
    # At least 2 edges ("knows", "works_with"), "has_secret" should be dropped
    assert stats["edge_count"] >= 2, f"Expected at least 2 edges, got {stats['edge_count']}"  # nosec B101

    # Verify dropped edges were counted
    assert memory._dropped_edges_count >= 1, "Should have dropped at least 1 edge"  # nosec B101

    print("✓ test_gsr_filters_not_erases PASSED")


if __name__ == "__main__":
    test_graph_sanity_integration()
    test_dual_retrieve_improves_recall()
    test_summarise_preserves_graph()
    test_concurrent_ingestion_warning()
    test_concurrent_read_write_stability()
    test_summarise_no_duplicate_edges()
    test_graph_is_persisted()
    test_gsr_filters_not_erases()
    print("\n✓ All Phase 3 + Invariant tests PASSED!")
