"""
test_conversational_edges.py
-----------------------------
Unit tests for conversational edge types: TURN_FOLLOWS, SPOKEN_BY, IN_SESSION.

These edges enable temporal/speaker/session-based queries in HIEROMEM.

Author: Antigravity
"""

import pytest
import numpy as np
import networkx as nx

from hieromem_core.memory import Memory
from hieromem_core.edge_types import (
    TURN_FOLLOWS, SPOKEN_BY, IN_SESSION,
    SPEAKER_NODE_PREFIX, SESSION_NODE_PREFIX,
    ConversationalEdgeBuilder,
    compute_temporal_weight,
    is_structural_edge,
)
from hieromem_core.persistence.in_memory_graph_store import InMemoryGraphStore


# ==============================================================================
# Edge Builder Tests
# ==============================================================================

class TestConversationalEdgeBuilder:
    """Tests for the ConversationalEdgeBuilder class."""

    def test_turn_follows_edge_basic(self):
        """Build a basic TURN_FOLLOWS edge."""
        edge = ConversationalEdgeBuilder.build_turn_follows_edge(
            prev_episode_id="ep_1",
            curr_episode_id="ep_2"
        )

        assert edge[0] == "ep_1"  # nosec B101
        assert edge[1] == "ep_2"  # nosec B101
        assert edge[2]["relation_type"] == TURN_FOLLOWS  # nosec B101
        assert edge[2]["weight"] == 1.0  # Adjacent turns = full weight  # nosec B101
        assert edge[2]["temporal_distance"] == 1  # nosec B101
        assert edge[2]["trust"] == 1.0  # nosec B101

    def test_turn_follows_edge_with_decay(self):
        """Weight decays with temporal distance."""
        edge = ConversationalEdgeBuilder.build_turn_follows_edge(
            prev_episode_id="ep_1",
            curr_episode_id="ep_3",
            temporal_distance=2,
            weight_decay=0.9
        )

        assert edge[2]["temporal_distance"] == 2  # nosec B101
        assert edge[2]["weight"] == pytest.approx(0.9, rel=1e-3)  # nosec B101

        # Distance 3 should decay more
        edge3 = ConversationalEdgeBuilder.build_turn_follows_edge(
            prev_episode_id="ep_1",
            curr_episode_id="ep_4",
            temporal_distance=3,
            weight_decay=0.9
        )
        assert edge3[2]["weight"] == pytest.approx(0.81, rel=1e-3)  # nosec B101

    def test_spoken_by_edge(self):
        """Build a SPOKEN_BY edge to speaker node."""
        edge = ConversationalEdgeBuilder.build_spoken_by_edge(
            episode_id="ep_1",
            speaker="Alice"
        )

        assert edge[0] == "ep_1"  # nosec B101
        assert edge[1] == f"{SPEAKER_NODE_PREFIX}Alice"  # nosec B101
        assert edge[2]["relation_type"] == SPOKEN_BY  # nosec B101
        assert edge[2]["speaker"] == "Alice"  # nosec B101
        assert edge[2]["weight"] == 1.0  # nosec B101

    def test_spoken_by_edge_strips_whitespace(self):
        """Speaker names are normalized."""
        edge = ConversationalEdgeBuilder.build_spoken_by_edge(
            episode_id="ep_1",
            speaker="  Bob  "
        )

        assert edge[1] == f"{SPEAKER_NODE_PREFIX}Bob"  # nosec B101
        assert edge[2]["speaker"] == "Bob"  # nosec B101

    def test_in_session_edge(self):
        """Build an IN_SESSION edge to session node."""
        edge = ConversationalEdgeBuilder.build_in_session_edge(
            episode_id="ep_1",
            session_id="session_1",
            turn_index=5,
            session_timestamp="2023_05_07"
        )

        assert edge[0] == "ep_1"  # nosec B101
        assert edge[1] == f"{SESSION_NODE_PREFIX}session_1"  # nosec B101
        assert edge[2]["relation_type"] == IN_SESSION  # nosec B101
        assert edge[2]["session_id"] == "session_1"  # nosec B101
        assert edge[2]["turn_index"] == 5  # nosec B101
        assert edge[2]["session_timestamp"] == "2023_05_07"  # nosec B101


class TestTemporalWeightComputation:
    """Tests for temporal weight decay computation."""

    def test_distance_1_full_weight(self):
        """Adjacent turns have full weight."""
        assert compute_temporal_weight(1) == 1.0  # nosec B101

    def test_decay_with_distance(self):
        """Weight decays exponentially with distance."""
        w2 = compute_temporal_weight(2, decay_rate=0.9)
        w3 = compute_temporal_weight(3, decay_rate=0.9)

        assert w2 == pytest.approx(0.9, rel=1e-3)  # nosec B101
        assert w3 == pytest.approx(0.81, rel=1e-3)  # nosec B101

    def test_min_weight_floor(self):
        """Weight doesn't drop below minimum."""
        w = compute_temporal_weight(100, decay_rate=0.5, min_weight=0.1)
        assert w >= 0.1  # nosec B101


class TestStructuralEdgeUtilities:
    """Tests for edge type utility functions."""

    def test_is_structural_edge(self):
        """Correctly identify structural vs semantic edges."""
        assert is_structural_edge(TURN_FOLLOWS) is True  # nosec B101
        assert is_structural_edge(SPOKEN_BY) is True  # nosec B101
        assert is_structural_edge(IN_SESSION) is True  # nosec B101
        assert is_structural_edge("related_to") is False  # nosec B101
        assert is_structural_edge("works_at") is False  # nosec B101


# ==============================================================================
# InMemoryGraphStore Integration Tests
# ==============================================================================

class TestInMemoryGraphStoreConversationalEdges:
    """Tests for conversational edges in InMemoryGraphStore."""

    @pytest.fixture
    def store(self):
        return InMemoryGraphStore()

    def test_add_structural_edge(self, store):
        """Add a structural edge to the graph."""
        edge = ConversationalEdgeBuilder.build_turn_follows_edge(
            "ep_1", "ep_2"
        )
        store.add_structural_edge(edge[0], edge[1], edge[2])

        # Verify edge exists
        assert store._G.number_of_edges() >= 1  # nosec B101

        # Verify we can query it
        result = store.query_turns_after("ep_1", limit=5)
        assert "ep_2" in result  # nosec B101

    def test_query_turns_after(self, store):
        """Query TURN_FOLLOWS chain."""
        # Build chain: ep_1 -> ep_2 -> ep_3
        for i in range(1, 3):
            edge = ConversationalEdgeBuilder.build_turn_follows_edge(
                f"ep_{i}", f"ep_{i+1}"
            )
            store.add_structural_edge(edge[0], edge[1], edge[2])

        # Query from ep_1
        result = store.query_turns_after("ep_1", limit=10)

        assert result == ["ep_2", "ep_3"]  # nosec B101

    def test_query_by_speaker(self, store):
        """Query episodes by speaker."""
        # Add episodes from different speakers
        for ep_id in ["ep_1", "ep_2", "ep_3"]:
            speaker = "Alice" if ep_id in ["ep_1", "ep_3"] else "Bob"
            edge = ConversationalEdgeBuilder.build_spoken_by_edge(ep_id, speaker)
            store.add_structural_edge(edge[0], edge[1], edge[2])

        # Query Alice's episodes
        alice_eps = store.query_by_speaker("Alice")
        assert len(alice_eps) == 2  # nosec B101
        assert "ep_1" in alice_eps  # nosec B101
        assert "ep_3" in alice_eps  # nosec B101

        # Query Bob's episodes
        bob_eps = store.query_by_speaker("Bob")
        assert bob_eps == ["ep_2"]  # nosec B101

    def test_query_by_session(self, store):
        """Query episodes by session."""
        # Add episodes to different sessions
        edge1 = ConversationalEdgeBuilder.build_in_session_edge(
            "ep_1", "session_A", turn_index=0
        )
        edge2 = ConversationalEdgeBuilder.build_in_session_edge(
            "ep_2", "session_A", turn_index=1
        )
        edge3 = ConversationalEdgeBuilder.build_in_session_edge(
            "ep_3", "session_B", turn_index=0
        )

        store.add_structural_edge(edge1[0], edge1[1], edge1[2])
        store.add_structural_edge(edge2[0], edge2[1], edge2[2])
        store.add_structural_edge(edge3[0], edge3[1], edge3[2])

        # Query session A
        session_a = store.query_by_session("session_A")
        assert len(session_a) == 2  # nosec B101
        assert session_a == ["ep_1", "ep_2"]  # Ordered by turn_index  # nosec B101

        # Query session B
        session_b = store.query_by_session("session_B")
        assert session_b == ["ep_3"]  # nosec B101

    def test_get_speakers(self, store):
        """Get all speaker names."""
        for speaker in ["Alice", "Bob", "Charlie"]:
            edge = ConversationalEdgeBuilder.build_spoken_by_edge(
                f"ep_{speaker}", speaker
            )
            store.add_structural_edge(edge[0], edge[1], edge[2])

        speakers = store.get_speakers()
        assert set(speakers) == {"Alice", "Bob", "Charlie"}  # nosec B101

    def test_get_sessions(self, store):
        """Get all session IDs."""
        for sess in ["session_1", "session_2"]:
            edge = ConversationalEdgeBuilder.build_in_session_edge(
                f"ep_{sess}", sess
            )
            store.add_structural_edge(edge[0], edge[1], edge[2])

        sessions = store.get_sessions()
        assert set(sessions) == {"session_1", "session_2"}  # nosec B101


# ==============================================================================
# Memory Integration Tests
# ==============================================================================

class TestMemoryConversationalEdges:
    """Tests for conversational edge creation in Memory class."""

    @pytest.fixture
    def memory(self):
        return Memory(
            auto_graph_extraction=False,
            debug=False,
            Nmax=100
        )

    def test_turn_follows_edge_created(self, memory):
        """TURN_FOLLOWS edges are created between consecutive episodes."""
        # Insert two episodes
        vec = np.random.randn(384).astype(np.float32)
        ep1 = memory.insert(vec, nx.DiGraph(), meta={"text": "first"})
        ep2 = memory.insert(vec, nx.DiGraph(), meta={"text": "second"})

        # Query following turns
        following = memory.get_following_turns(ep1.id, k=5)

        assert len(following) == 1  # nosec B101
        assert following[0].id == ep2.id  # nosec B101

    def test_spoken_by_edge_created_with_speaker_meta(self, memory):
        """SPOKEN_BY edges created when speaker in metadata."""
        vec = np.random.randn(384).astype(np.float32)

        # Insert with speaker metadata
        ep = memory.insert(vec, nx.DiGraph(), meta={"speaker": "Alice"})

        # Query by speaker
        alice_eps = memory.get_by_speaker("Alice")

        assert len(alice_eps) >= 1  # nosec B101
        assert ep.id in [e.id for e in alice_eps]  # nosec B101

    def test_in_session_edge_created_with_session_meta(self, memory):
        """IN_SESSION edges created when session_id in metadata."""
        vec = np.random.randn(384).astype(np.float32)

        # Insert with session metadata
        ep = memory.insert(
            vec, nx.DiGraph(),
            meta={"session_id": "session_1", "turn_idx": 0}
        )

        # Query by session
        session_eps = memory.get_in_session("session_1")

        assert len(session_eps) >= 1  # nosec B101
        assert ep.id in [e.id for e in session_eps]  # nosec B101

    def test_reset_conversation_breaks_turn_chain(self, memory):
        """reset_conversation() prevents TURN_FOLLOWS between unrelated episodes."""
        vec = np.random.randn(384).astype(np.float32)

        # First conversation
        memory.insert(vec, nx.DiGraph(), meta={"text": "conv1_turn1"})
        ep2 = memory.insert(vec, nx.DiGraph(), meta={"text": "conv1_turn2"})

        # Reset for new conversation
        memory.reset_conversation()

        # Second conversation (unrelated)
        ep3 = memory.insert(vec, nx.DiGraph(), meta={"text": "conv2_turn1"})

        # ep2 should NOT link to ep3
        following = memory.get_following_turns(ep2.id, k=5)

        # ep3 should not be in the chain from ep2
        assert ep3.id not in [e.id for e in following]  # nosec B101

    def test_multiple_speakers_in_chunk(self, memory):
        """Handle list of speakers in metadata."""
        vec = np.random.randn(384).astype(np.float32)

        # Insert with multiple speakers
        ep = memory.insert(
            vec, nx.DiGraph(),
            meta={"speakers": ["Alice", "Bob"]}
        )

        # Both should find this episode
        alice_eps = memory.get_by_speaker("Alice")
        bob_eps = memory.get_by_speaker("Bob")

        assert ep.id in [e.id for e in alice_eps]  # nosec B101
        assert ep.id in [e.id for e in bob_eps]  # nosec B101


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
