"""Integration tests for TraceMem with real Neo4j and LanceDB.

These tests verify end-to-end data persistence:
- Messages are correctly stored in Neo4j as graph nodes
- User text is indexed in LanceDB for vector search
- Traces with tool calls create proper resource nodes and edges

Start Neo4j with:
    docker compose up -d neo4j

Run tests with:
    uv run pytest tests/integration/test_tracemem_integration.py -v
"""

import hashlib

import pytest

from tracemem_core.extractors import _canonicalize_file_uri
from tracemem_core.messages import Message, ToolCall
from tracemem_core.retrieval import HybridRetrievalStrategy, RetrievalConfig
from tracemem_core.tracemem import TraceMem

pytestmark = pytest.mark.neo4j


class TestUserMessageIngestion:
    """Tests for user message ingestion into Neo4j and LanceDB."""

    async def test_add_user_message_creates_neo4j_node(
        self, tracemem_integration: TraceMem
    ):
        """Verify user message creates UserText node in Neo4j."""
        message = Message(role="user", content="Hello world")
        result = await tracemem_integration.add_message("conv-1", message)

        assert "user_text" in result
        user_text_id = result["user_text"]

        # Verify node exists in Neo4j
        user_text = await tracemem_integration._graph_store.get_user_text(user_text_id)
        assert user_text is not None
        assert user_text.text == "Hello world"
        assert user_text.conversation_id == "conv-1"

    async def test_add_user_message_indexes_in_lancedb(
        self, tracemem_integration: TraceMem
    ):
        """Verify user message is indexed in LanceDB."""
        message = Message(role="user", content="Search test content")
        result = await tracemem_integration.add_message("conv-1", message)

        user_text_id = result["user_text"]

        # Search should find the message
        query_vector = await tracemem_integration._embedder.embed("Search test")
        search_results = await tracemem_integration._vector_store.search(
            query_vector=query_vector,
            query_text="Search test",
            limit=10,
        )
        assert len(search_results) >= 1
        assert any(r.node_id == user_text_id for r in search_results)

    async def test_add_multiple_user_messages_all_searchable(
        self, tracemem_integration: TraceMem
    ):
        """Verify multiple user messages are all indexed and searchable."""
        messages = [
            ("conv-1", "First message about Python programming"),
            ("conv-1", "Second message about database queries"),
            ("conv-2", "Third message about API design"),
        ]

        created_ids = []
        for conv_id, content in messages:
            result = await tracemem_integration.add_message(
                conv_id, Message(role="user", content=content)
            )
            created_ids.append(result["user_text"])

        # Search for programming-related content
        query_vector = await tracemem_integration._embedder.embed("programming")
        results = await tracemem_integration._vector_store.search(
            query_vector=query_vector,
            query_text="programming",
            limit=10,
        )

        # All messages should be in the store
        found_ids = {r.node_id for r in results}
        assert len(found_ids) == 3


class TestAssistantMessageIngestion:
    """Tests for assistant message ingestion and graph relationships."""

    async def test_add_assistant_message_creates_agent_text_node(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify assistant message creates AgentText node in Neo4j."""
        # Add user first
        user_result = await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Hi")
        )
        user_text_id = user_result["user_text"]

        # Add assistant
        result = await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Hello!")
        )

        assert "agent_text" in result

        # Verify context shows both using retrieval strategy
        context = await retrieval_strategy.get_context(user_text_id)
        assert context.user_text is not None
        assert context.user_text.text == "Hi"
        assert context.agent_text is not None
        assert context.agent_text.text == "Hello!"

    async def test_assistant_message_links_to_user_message(
        self, tracemem_integration: TraceMem
    ):
        """Verify MESSAGE edge connects UserText to AgentText."""
        user_result = await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="What is Python?")
        )
        user_text_id = user_result["user_text"]

        await tracemem_integration.add_message(
            "conv-1",
            Message(role="assistant", content="Python is a programming language."),
        )

        # Query graph for the relationship
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (u:UserText {id: $id})-[r:MESSAGE]->(a:AgentText)
                RETURN u.text as user_text, a.text as agent_text
                """,
                id=str(user_text_id),
            )
            record = await result.single()
            assert record is not None
            assert record["user_text"] == "What is Python?"
            assert record["agent_text"] == "Python is a programming language."


class TestFullTraceImport:
    """Tests for importing complete traces with tool calls."""

    async def test_import_trace_creates_full_graph(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify import_trace creates complete graph structure."""
        test_file = tmp_path / "test.py"
        test_file.write_text("def hello(): pass")

        messages = [
            Message(role="user", content="Read test.py"),
            Message(
                role="assistant",
                content="Reading the file...",
                tool_calls=[
                    ToolCall(id="call_1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="def hello(): pass", tool_call_id="call_1"),
        ]

        result = await tracemem_integration.import_trace("conv-1", messages)

        # Verify all nodes created
        assert "user_text" in result
        assert "agent_text" in result

        # Verify Resource and ResourceVersion created
        resource_keys = [k for k in result.keys() if k.startswith("resource_") and not k.startswith("resource_version_")]
        version_keys = [k for k in result.keys() if k.startswith("resource_version_")]

        assert len(resource_keys) == 1, f"Expected 1 resource, got {resource_keys}"
        assert len(version_keys) == 1, f"Expected 1 version, got {version_keys}"

        # Verify graph traversal works using retrieval strategy
        context = await retrieval_strategy.get_context(result["user_text"])
        assert context.user_text.text == "Read test.py"
        assert context.agent_text.text == "Reading the file..."
        assert len(context.tool_uses) == 1

    async def test_import_trace_with_multiple_tool_calls(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify import_trace handles multiple tool calls in one message."""
        file1 = tmp_path / "auth.py"
        file2 = tmp_path / "utils.py"
        file1.write_text("def login(): pass")
        file2.write_text("def helper(): pass")

        messages = [
            Message(role="user", content="Read auth.py and utils.py"),
            Message(
                role="assistant",
                content="Reading both files...",
                tool_calls=[
                    ToolCall(id="call_1", name="read_file", args={"path": str(file1)}),
                    ToolCall(id="call_2", name="read_file", args={"path": str(file2)}),
                ],
            ),
            Message(role="tool", content="def login(): pass", tool_call_id="call_1"),
            Message(role="tool", content="def helper(): pass", tool_call_id="call_2"),
        ]

        result = await tracemem_integration.import_trace("conv-1", messages)

        # Should have 2 resources and 2 versions
        resource_keys = [k for k in result.keys() if k.startswith("resource_") and not k.startswith("resource_version_")]
        version_keys = [k for k in result.keys() if k.startswith("resource_version_")]

        assert len(resource_keys) == 2
        assert len(version_keys) == 2

        # Verify tool uses in context using retrieval strategy
        context = await retrieval_strategy.get_context(result["user_text"])
        assert len(context.tool_uses) == 2


class TestResourceVersioning:
    """Tests for resource versioning when content changes."""

    async def test_resource_versioning_on_content_change(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify new ResourceVersion created when content changes."""
        test_file = tmp_path / "code.py"

        # First version
        messages1 = [
            Message(role="user", content="Read code.py"),
            Message(
                role="assistant",
                content="Reading...",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="v1", tool_call_id="c1"),
        ]
        await tracemem_integration.import_trace("conv-1", messages1)

        # Second version (different content)
        messages2 = [
            Message(role="user", content="Read code.py again"),
            Message(
                role="assistant",
                content="Reading again...",
                tool_calls=[
                    ToolCall(id="c2", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="v2", tool_call_id="c2"),
        ]
        await tracemem_integration.import_trace("conv-2", messages2)

        # Verify single Resource with updated hash
        canonical_uri = _canonicalize_file_uri(f"file://{test_file}", root=None)
        resource = await tracemem_integration._graph_store.get_resource_by_uri(
            canonical_uri
        )
        assert resource is not None

        # Hash should be for v2
        expected_hash = hashlib.sha256("v2".encode()).hexdigest()
        assert resource.current_content_hash == expected_hash

    async def test_same_content_does_not_create_new_version(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify same content hash does not create duplicate versions."""
        test_file = tmp_path / "stable.py"

        # First read
        messages1 = [
            Message(role="user", content="Read stable.py"),
            Message(
                role="assistant",
                content="Reading...",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="same content", tool_call_id="c1"),
        ]
        result1 = await tracemem_integration.import_trace("conv-1", messages1)

        # Second read with same content
        messages2 = [
            Message(role="user", content="Read stable.py again"),
            Message(
                role="assistant",
                content="Reading again...",
                tool_calls=[
                    ToolCall(id="c2", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="same content", tool_call_id="c2"),
        ]
        result2 = await tracemem_integration.import_trace("conv-2", messages2)

        # First import should create resource and version
        version_keys_1 = [k for k in result1.keys() if k.startswith("resource_version_")]
        assert len(version_keys_1) == 1

        # Second import should NOT create new version (same hash)
        version_keys_2 = [k for k in result2.keys() if k.startswith("resource_version_")]
        assert len(version_keys_2) == 0, "Same content should not create new version"

    async def test_version_of_edges_created(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify VERSION_OF edges link ResourceVersion to Resource."""
        test_file = tmp_path / "versioned.py"

        messages = [
            Message(role="user", content="Read versioned.py"),
            Message(
                role="assistant",
                content="Reading...",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="version 1", tool_call_id="c1"),
        ]
        await tracemem_integration.import_trace("conv-1", messages)

        # Query for VERSION_OF relationship
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (v:ResourceVersion)-[r:VERSION_OF]->(res:Resource)
                RETURN v.content_hash as version_hash, res.uri as resource_uri
                """
            )
            records = await result.data()
            assert len(records) == 1
            assert records[0]["version_hash"] == hashlib.sha256("version 1".encode()).hexdigest()


class TestMultipleConversations:
    """Tests for conversation isolation and cross-conversation queries."""

    async def test_multiple_conversations_isolated(
        self, tracemem_integration: TraceMem
    ):
        """Verify conversations are properly isolated in search."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Message from conversation 1")
        )
        await tracemem_integration.add_message(
            "conv-2", Message(role="user", content="Message from conversation 2")
        )

        # Search excluding conv-1 should not return conv-1 content
        query_vector = await tracemem_integration._embedder.embed("Message")
        results = await tracemem_integration._vector_store.search(
            query_vector=query_vector,
            query_text="Message",
            limit=10,
            exclude_conversation_id="conv-1",
        )

        assert len(results) >= 1
        assert all(r.conversation_id != "conv-1" for r in results)

    async def test_conversations_share_resources(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify multiple conversations accessing same file share Resource node."""
        shared_file = tmp_path / "shared.py"

        # First conversation reads the file
        messages1 = [
            Message(role="user", content="Read shared.py"),
            Message(
                role="assistant",
                content="Reading...",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(shared_file)})
                ],
            ),
            Message(role="tool", content="shared content", tool_call_id="c1"),
        ]
        result1 = await tracemem_integration.import_trace("conv-1", messages1)

        # Second conversation reads the same file
        messages2 = [
            Message(role="user", content="Read shared.py too"),
            Message(
                role="assistant",
                content="Reading again...",
                tool_calls=[
                    ToolCall(id="c2", name="read_file", args={"path": str(shared_file)})
                ],
            ),
            Message(role="tool", content="shared content", tool_call_id="c2"),
        ]
        result2 = await tracemem_integration.import_trace("conv-2", messages2)

        # Both should reference the same Resource (MERGE behavior)
        resource_keys_1 = [k for k in result1.keys() if k.startswith("resource_") and not k.startswith("resource_version_")]
        resource_keys_2 = [k for k in result2.keys() if k.startswith("resource_") and not k.startswith("resource_version_")]

        # First conversation creates the resource
        assert len(resource_keys_1) == 1

        # Second conversation should NOT create a new resource
        assert len(resource_keys_2) == 0, "Second conversation should reuse existing resource"

    async def test_get_conversations_for_resource(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify we can find all conversations that touched a resource."""
        shared_file = tmp_path / "popular.py"

        # Multiple conversations read the same file
        for i in range(3):
            messages = [
                Message(role="user", content=f"Read popular.py - conv {i}"),
                Message(
                    role="assistant",
                    content="Reading...",
                    tool_calls=[
                        ToolCall(
                            id=f"c{i}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=f"content v{i}", tool_call_id=f"c{i}"),
            ]
            await tracemem_integration.import_trace(f"conv-{i}", messages)

        # Query for all conversations that touched this resource using strategy
        canonical_uri = _canonicalize_file_uri(f"file://{shared_file}", root=None)
        conversations = await retrieval_strategy.get_conversations_for_resource(
            canonical_uri
        )

        assert len(conversations) == 3
        conv_ids = {c.conversation_id for c in conversations}
        assert conv_ids == {"conv-0", "conv-1", "conv-2"}


class TestCrossConversationSharedResource:
    """Tests for cross-conversation queries when multiple conversations access the same file.

    This test class verifies the notebook example scenario:
    - Two conversations read the same file with identical content
    - Both create READ_FILE edges to the same ResourceVersion (no new version created)
    - Cross-conversation query finds BOTH conversations that touched the file
    """

    async def test_two_conversations_same_file_same_content(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Two conversations reading same file with same content should both be queryable.

        This is the key scenario from the notebook:
        - Conversation 1 reads config.py -> creates Resource, ResourceVersion, READ_FILE edge
        - Conversation 2 reads config.py (same content) -> reuses ResourceVersion, creates READ_FILE edge
        - Query for "conversations that touched config.py" should return BOTH conversations
        """
        config_file = tmp_path / "config.py"
        file_content = "API_KEY = 'secret'\nDEBUG = True"

        # Conversation 1: First developer reads config.py
        messages1 = [
            Message(role="user", content="Show me the config file"),
            Message(
                role="assistant",
                content="I'll read the configuration file for you.",
                tool_calls=[
                    ToolCall(
                        id="call_1",
                        name="read_file",
                        args={"path": str(config_file)},
                    )
                ],
            ),
            Message(role="tool", content=file_content, tool_call_id="call_1"),
            Message(
                role="assistant",
                content="Here's your config - it has an API key and debug flag.",
            ),
        ]
        result1 = await tracemem_integration.import_trace("conv-alice", messages1)

        # Conversation 2: Second developer reads the same file (same content)
        messages2 = [
            Message(role="user", content="What's in the config?"),
            Message(
                role="assistant",
                content="Let me check the configuration.",
                tool_calls=[
                    ToolCall(
                        id="call_2",
                        name="read_file",
                        args={"path": str(config_file)},
                    )
                ],
            ),
            Message(role="tool", content=file_content, tool_call_id="call_2"),
            Message(
                role="assistant",
                content="The config contains API_KEY and DEBUG settings.",
            ),
        ]
        result2 = await tracemem_integration.import_trace("conv-bob", messages2)

        # Verify Resource was created only once
        resource_keys_1 = [
            k for k in result1.keys()
            if k.startswith("resource_") and not k.startswith("resource_version_")
        ]
        resource_keys_2 = [
            k for k in result2.keys()
            if k.startswith("resource_") and not k.startswith("resource_version_")
        ]
        assert len(resource_keys_1) == 1, "First conv should create Resource"
        assert len(resource_keys_2) == 0, "Second conv should reuse Resource"

        # Verify ResourceVersion was created only once (same content hash)
        version_keys_1 = [k for k in result1.keys() if k.startswith("resource_version_")]
        version_keys_2 = [k for k in result2.keys() if k.startswith("resource_version_")]
        assert len(version_keys_1) == 1, "First conv should create ResourceVersion"
        assert len(version_keys_2) == 0, "Second conv should reuse ResourceVersion (same hash)"

        # THE KEY TEST: Both conversations should be found via cross-conversation query
        canonical_uri = _canonicalize_file_uri(f"file://{config_file}", root=None)
        conversations = await retrieval_strategy.get_conversations_for_resource(canonical_uri)

        conv_ids = {c.conversation_id for c in conversations}
        assert conv_ids == {"conv-alice", "conv-bob"}, (
            f"Expected both conversations to be found, got {conv_ids}"
        )

    async def test_shared_resource_graph_structure(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify the graph structure when two conversations share a resource.

        Expected graph:
        conv-alice:
            UserText1 -[MESSAGE]-> AgentText1 -[READ_FILE]-> ResourceVersion <-[VERSION_OF]- Resource
                                             -[MESSAGE]-> AgentText2

        conv-bob:
            UserText2 -[MESSAGE]-> AgentText3 -[READ_FILE]-> ResourceVersion (same node!)
                                             -[MESSAGE]-> AgentText4
        """
        shared_file = tmp_path / "shared.py"
        content = "# shared code"

        # Two conversations read the same file
        for conv_id in ["conv-alice", "conv-bob"]:
            messages = [
                Message(role="user", content=f"Read shared.py ({conv_id})"),
                Message(
                    role="assistant",
                    content="Reading...",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{conv_id}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{conv_id}"),
            ]
            await tracemem_integration.import_trace(conv_id, messages)

        # Query the graph to verify structure
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            # Verify there's only ONE ResourceVersion node
            result = await session.run(
                """
                MATCH (v:ResourceVersion)
                RETURN count(v) as version_count
                """
            )
            record = await result.single()
            assert record["version_count"] == 1, "Should have exactly one ResourceVersion"

            # Verify there's only ONE Resource node
            result = await session.run(
                """
                MATCH (r:Resource)
                RETURN count(r) as resource_count
                """
            )
            record = await result.single()
            assert record["resource_count"] == 1, "Should have exactly one Resource"

            # Verify TWO READ_FILE edges point to the same ResourceVersion
            result = await session.run(
                """
                MATCH (a:AgentText)-[r:READ_FILE]->(v:ResourceVersion)
                RETURN count(r) as edge_count, collect(DISTINCT a.conversation_id) as conversations
                """
            )
            record = await result.single()
            assert record["edge_count"] == 2, "Should have two READ_FILE edges"
            assert set(record["conversations"]) == {"conv-alice", "conv-bob"}

    async def test_cross_conversation_query_finds_all_users(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify cross-conversation query returns user text from each conversation."""
        auth_file = tmp_path / "auth.py"
        content = "def login(user, password): pass"

        # Three different conversations read the same auth file
        user_messages = [
            ("conv-1", "How does authentication work?"),
            ("conv-2", "Show me the login code"),
            ("conv-3", "Review the auth implementation"),
        ]

        for conv_id, user_msg in user_messages:
            messages = [
                Message(role="user", content=user_msg),
                Message(
                    role="assistant",
                    content="Let me read that file.",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{conv_id}",
                            name="read_file",
                            args={"path": str(auth_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{conv_id}"),
            ]
            await tracemem_integration.import_trace(conv_id, messages)

        # Query for all conversations that touched auth.py
        canonical_uri = _canonicalize_file_uri(f"file://{auth_file}", root=None)
        conversations = await retrieval_strategy.get_conversations_for_resource(canonical_uri)

        assert len(conversations) == 3

        # Verify each conversation reference includes the user's original question
        user_texts = {c.user_text for c in conversations}
        expected_texts = {msg for _, msg in user_messages}
        assert user_texts == expected_texts

    async def test_resource_version_lookup_by_hash(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify get_resource_version_by_hash returns the correct version."""
        test_file = tmp_path / "test.py"
        content = "def test(): pass"
        content_hash = hashlib.sha256(content.encode()).hexdigest()

        # Create a conversation that reads the file
        messages = [
            Message(role="user", content="Read test.py"),
            Message(
                role="assistant",
                content="Reading...",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content=content, tool_call_id="c1"),
        ]
        await tracemem_integration.import_trace("conv-1", messages)

        # Look up the version by URI and hash
        canonical_uri = _canonicalize_file_uri(f"file://{test_file}", root=None)
        version = await tracemem_integration._graph_store.get_resource_version_by_hash(
            canonical_uri, content_hash
        )

        assert version is not None
        assert version.uri == canonical_uri
        assert version.content_hash == content_hash

        # Non-existent hash should return None
        no_version = await tracemem_integration._graph_store.get_resource_version_by_hash(
            canonical_uri, "nonexistent_hash"
        )
        assert no_version is None


class TestConversationChain:
    """Tests for conversation chain continuity (fully connected subgraphs)."""

    async def test_conversation_forms_connected_chain(
        self, tracemem_integration: TraceMem
    ):
        """Verify alternating messages form a fully connected chain.

        Graph structure should be:
        U1 -> A1 -> U2 -> A2 -> U3 -> A3
        """
        # Add multiple user/assistant message pairs
        for i in range(3):
            await tracemem_integration.add_message(
                "conv-1", Message(role="user", content=f"User message {i + 1}")
            )
            await tracemem_integration.add_message(
                "conv-1", Message(role="assistant", content=f"Agent response {i + 1}")
            )

        # Query the graph to verify the chain is connected
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            # Find the chain length by following all MESSAGE edges
            result = await session.run(
                """
                MATCH path = (start)-[:MESSAGE*]->(end)
                WHERE start.conversation_id = 'conv-1'
                  AND NOT ()-[:MESSAGE]->(start)
                RETURN length(path) as chain_length
                ORDER BY chain_length DESC
                LIMIT 1
                """
            )
            record = await result.single()
            # 3 user + 3 agent = 6 nodes, so 5 edges in chain
            assert record is not None
            assert record["chain_length"] == 5

    async def test_new_user_message_links_from_last_agent(
        self, tracemem_integration: TraceMem
    ):
        """Verify a new user message creates edge from last agent node."""
        # Add initial pair
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="First user message")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="First agent response")
        )

        # Add second user message - should link FROM agent TO this user
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Second user message")
        )

        # Verify the agent->user edge exists
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (a:AgentText {text: 'First agent response'})
                      -[:MESSAGE]->
                      (u:UserText {text: 'Second user message'})
                RETURN a.text as agent_text, u.text as user_text
                """
            )
            record = await result.single()
            assert record is not None
            assert record["agent_text"] == "First agent response"
            assert record["user_text"] == "Second user message"

    async def test_conversation_chain_uses_timestamp_ordering(
        self, tracemem_integration: TraceMem
    ):
        """Verify chain links use timestamp-based node lookup, not memory."""
        # Create first pair
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="User 1")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Agent 1")
        )

        # Simulate a "new session" by clearing any hypothetical memory state
        # (this is now a no-op since we removed _last_user_text)

        # Add more messages - they should still link correctly via timestamp
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="User 2")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Agent 2")
        )

        # Verify the full chain exists
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (u1:UserText {text: 'User 1'})
                      -[:MESSAGE]->(a1:AgentText {text: 'Agent 1'})
                      -[:MESSAGE]->(u2:UserText {text: 'User 2'})
                      -[:MESSAGE]->(a2:AgentText {text: 'Agent 2'})
                RETURN count(*) as chain_exists
                """
            )
            record = await result.single()
            assert record is not None
            assert record["chain_exists"] == 1

    async def test_tool_usage_creates_connected_chain(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify tool usage with multiple assistant messages forms connected chain.

        Flow: User -> Assistant (with tool_call) -> Tool result -> Assistant (final)
        Graph: UserText -> AgentText1 -> AgentText2
        """
        test_file = tmp_path / "code.py"
        test_file.write_text("print('hello')")

        messages = [
            Message(role="user", content="Read code.py and explain it"),
            Message(
                role="assistant",
                content="I'll read the file for you.",
                tool_calls=[
                    ToolCall(id="call_1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="print('hello')", tool_call_id="call_1"),
            Message(
                role="assistant",
                content="The file contains a simple print statement.",
            ),
        ]

        await tracemem_integration.import_trace("conv-1", messages)

        # Verify the chain: UserText -> AgentText1 -> AgentText2
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (u:UserText {text: 'Read code.py and explain it'})
                      -[:MESSAGE]->(a1:AgentText {text: "I'll read the file for you."})
                      -[:MESSAGE]->(a2:AgentText {text: 'The file contains a simple print statement.'})
                RETURN count(*) as chain_exists
                """
            )
            record = await result.single()
            assert record is not None
            assert record["chain_exists"] == 1, "Tool usage should create AgentText -> AgentText chain"

    async def test_multiple_tool_calls_with_final_response(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify multiple tool calls followed by final response creates proper chain."""
        file1 = tmp_path / "a.py"
        file2 = tmp_path / "b.py"
        file1.write_text("# file a")
        file2.write_text("# file b")

        messages = [
            Message(role="user", content="Compare a.py and b.py"),
            Message(
                role="assistant",
                content="Let me read both files.",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(file1)}),
                    ToolCall(id="c2", name="read_file", args={"path": str(file2)}),
                ],
            ),
            Message(role="tool", content="# file a", tool_call_id="c1"),
            Message(role="tool", content="# file b", tool_call_id="c2"),
            Message(role="assistant", content="Both files are similar comment headers."),
        ]

        await tracemem_integration.import_trace("conv-1", messages)

        # Verify chain length: User -> Agent1 -> Agent2 = 2 edges
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH path = (u:UserText)-[:MESSAGE*]->(final:AgentText)
                WHERE u.conversation_id = 'conv-1'
                  AND u.text = 'Compare a.py and b.py'
                  AND NOT (final)-[:MESSAGE]->()
                RETURN length(path) as chain_length, final.text as final_text
                """
            )
            record = await result.single()
            assert record is not None
            assert record["chain_length"] == 2
            assert record["final_text"] == "Both files are similar comment headers."


class TestTurnIndex:
    """Tests for turn_index tracking in conversation nodes."""

    async def test_user_message_starts_new_turn(
        self, tracemem_integration: TraceMem
    ):
        """First user message should be turn 0."""
        result = await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Hello")
        )

        user_text = await tracemem_integration._graph_store.get_user_text(
            result["user_text"]
        )
        assert user_text is not None
        assert user_text.turn_index == 0

    async def test_assistant_same_turn_as_user(
        self, tracemem_integration: TraceMem
    ):
        """Assistant gets same turn_index as preceding user."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Hello")
        )
        result = await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Hi!")
        )

        # Query for the AgentText node
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            query_result = await session.run(
                """
                MATCH (a:AgentText {id: $id})
                RETURN a.turn_index as turn_index
                """,
                id=str(result["agent_text"]),
            )
            record = await query_result.single()
            assert record is not None
            assert record["turn_index"] == 0

    async def test_multiple_assistants_same_turn(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Tool flow keeps all assistants in same turn."""
        test_file = tmp_path / "code.py"
        test_file.write_text("print('hello')")

        messages = [
            Message(role="user", content="Read code.py and explain it"),
            Message(
                role="assistant",
                content="I'll read the file for you.",
                tool_calls=[
                    ToolCall(id="call_1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="print('hello')", tool_call_id="call_1"),
            Message(
                role="assistant",
                content="The file contains a simple print statement.",
            ),
        ]

        await tracemem_integration.import_trace("conv-1", messages)

        # Query for all AgentText nodes in the conversation
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (a:AgentText {conversation_id: 'conv-1'})
                RETURN a.turn_index as turn_index
                ORDER BY a.created_at
                """
            )
            records = await result.data()
            # Both assistant messages should be turn 0
            assert len(records) == 2
            assert all(r["turn_index"] == 0 for r in records)

    async def test_second_user_increments_turn(
        self, tracemem_integration: TraceMem
    ):
        """Second user message should be turn 1."""
        # Turn 0
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="User 1")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Agent 1")
        )

        # Turn 1
        result = await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="User 2")
        )

        user_text = await tracemem_integration._graph_store.get_user_text(
            result["user_text"]
        )
        assert user_text is not None
        assert user_text.turn_index == 1

    async def test_turn_based_chain_query(
        self, tracemem_integration: TraceMem
    ):
        """Can query all nodes in a specific turn."""
        # Create multiple turns
        for i in range(3):
            await tracemem_integration.add_message(
                "conv-1", Message(role="user", content=f"User {i}")
            )
            await tracemem_integration.add_message(
                "conv-1", Message(role="assistant", content=f"Agent {i}")
            )

        # Query nodes in turn 1
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (n)
                WHERE (n:UserText OR n:AgentText)
                  AND n.conversation_id = 'conv-1'
                  AND n.turn_index = 1
                RETURN n.text as text
                ORDER BY n.created_at
                """
            )
            records = await result.data()
            # Turn 1 should have User 1 and Agent 1
            assert len(records) == 2
            assert records[0]["text"] == "User 1"
            assert records[1]["text"] == "Agent 1"

    async def test_get_max_turn_index(
        self, tracemem_integration: TraceMem
    ):
        """Verify get_max_turn_index returns correct value."""
        # Empty conversation
        max_turn = await tracemem_integration._graph_store.get_max_turn_index(
            "conv-empty"
        )
        assert max_turn == -1

        # Add some turns
        for i in range(3):
            await tracemem_integration.add_message(
                "conv-1", Message(role="user", content=f"User {i}")
            )
            await tracemem_integration.add_message(
                "conv-1", Message(role="assistant", content=f"Agent {i}")
            )

        max_turn = await tracemem_integration._graph_store.get_max_turn_index("conv-1")
        assert max_turn == 2  # Turns are 0, 1, 2

    async def test_get_last_node_in_turn(
        self, tracemem_integration: TraceMem
    ):
        """Verify get_last_node_in_turn returns correct node."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="User message")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Agent message")
        )

        # Last node in turn 0 should be the agent
        last_node = await tracemem_integration._graph_store.get_last_node_in_turn(
            "conv-1", 0
        )
        assert last_node is not None
        assert last_node.text == "Agent message"


class TestToolUsesIntegration:
    """Tests for tool_uses property persistence in Neo4j."""

    async def test_tool_uses_persisted_in_neo4j(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify tool_uses are correctly stored and retrieved from Neo4j."""
        test_file = tmp_path / "code.py"
        test_file.write_text("print('hello')")

        messages = [
            Message(role="user", content="Check git status and read the file"),
            Message(
                role="assistant",
                content="I'll check git and read the file for you.",
                tool_calls=[
                    ToolCall(id="call_1", name="bash", args={"command": "git status"}),
                    ToolCall(id="call_2", name="read_file", args={"path": str(test_file)}),
                ],
            ),
            Message(role="tool", content="On branch main", tool_call_id="call_1"),
            Message(role="tool", content="print('hello')", tool_call_id="call_2"),
        ]

        await tracemem_integration.import_trace("conv-1", messages)

        # Query Neo4j directly to verify tool_uses storage
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (a:AgentText {conversation_id: 'conv-1'})
                RETURN a.tool_uses as tool_uses
                """
            )
            record = await result.single()
            assert record is not None

            import json
            tool_uses = json.loads(record["tool_uses"])
            assert len(tool_uses) == 2
            assert tool_uses[0]["id"] == "call_1"
            assert tool_uses[0]["name"] == "bash"
            assert tool_uses[0]["args"] == {"command": "git status"}
            assert tool_uses[1]["id"] == "call_2"
            assert tool_uses[1]["name"] == "read_file"

    async def test_tool_uses_empty_json_when_no_tools(
        self, tracemem_integration: TraceMem
    ):
        """Verify empty tool_uses is stored as empty JSON array."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Hello")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Hi there!")
        )

        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (a:AgentText {conversation_id: 'conv-1'})
                RETURN a.tool_uses as tool_uses
                """
            )
            record = await result.single()
            assert record is not None

            import json
            tool_uses = json.loads(record["tool_uses"])
            assert tool_uses == []

    async def test_non_resource_tools_tracked_in_tool_uses(
        self, tracemem_integration: TraceMem
    ):
        """Verify tools that don't create resources are still tracked in tool_uses."""
        messages = [
            Message(role="user", content="Run a command"),
            Message(
                role="assistant",
                content="Let me run that command.",
                tool_calls=[
                    # bash and web_search don't extract resources (no file_path/path arg)
                    ToolCall(id="c1", name="bash", args={"command": "echo hello"}),
                    ToolCall(id="c2", name="web_search", args={"query": "python docs"}),
                ],
            ),
            Message(role="tool", content="hello", tool_call_id="c1"),
            Message(role="tool", content="Search results...", tool_call_id="c2"),
        ]

        result = await tracemem_integration.import_trace("conv-1", messages)

        # These tools don't create resources, so no resource_* keys
        resource_keys = [k for k in result.keys() if k.startswith("resource")]
        assert len(resource_keys) == 0

        # But they should be tracked in tool_uses
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            query_result = await session.run(
                """
                MATCH (a:AgentText {conversation_id: 'conv-1'})
                RETURN a.tool_uses as tool_uses
                """
            )
            record = await query_result.single()

            import json
            tool_uses = json.loads(record["tool_uses"])
            assert len(tool_uses) == 2
            tool_names = [tu["name"] for tu in tool_uses]
            assert tool_names == ["bash", "web_search"]


class TestCrossConversationResourceEdges:
    """Tests verifying that multiple conversations create proper edges to shared resources.

    When two conversations reference the same resource:
    1. Both should have READ_FILE edges pointing to the same ResourceVersion
    2. Both conversation subgraphs should be connected via the shared ResourceVersion
    3. This enables cross-conversation queries to find all conversations that touched a file
    """

    async def test_two_conversations_both_create_edges_to_shared_resource(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Two conversations reading the same file should both have edges to ResourceVersion.

        Graph structure:
        conv-1: UserText1 -> AgentText1 -[READ_FILE]-> ResourceVersion <- Resource
        conv-2: UserText2 -> AgentText2 -[READ_FILE]-> ResourceVersion (SAME NODE)

        Both conversations should have paths to the same ResourceVersion node.
        """
        auth_file = tmp_path / "auth.py"
        content = "def login(user, pwd): pass"

        # Conversation 1 reads auth.py
        messages1 = [
            Message(role="user", content="Show me the auth code"),
            Message(
                role="assistant",
                content="Reading auth.py...",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(auth_file)})
                ],
            ),
            Message(role="tool", content=content, tool_call_id="c1"),
        ]
        await tracemem_integration.import_trace("conv-1", messages1)

        # Conversation 2 also reads auth.py (same content)
        messages2 = [
            Message(role="user", content="What's in the auth file?"),
            Message(
                role="assistant",
                content="Let me check auth.py...",
                tool_calls=[
                    ToolCall(id="c2", name="read_file", args={"path": str(auth_file)})
                ],
            ),
            Message(role="tool", content=content, tool_call_id="c2"),
        ]
        await tracemem_integration.import_trace("conv-2", messages2)

        # Verify both conversations have READ_FILE edges to the SAME ResourceVersion
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            # Count total READ_FILE edges - should be exactly 2
            result = await session.run(
                """
                MATCH (a:AgentText)-[r:READ_FILE]->(v:ResourceVersion)
                RETURN count(r) as edge_count
                """
            )
            record = await result.single()
            assert record["edge_count"] == 2, "Both conversations should have READ_FILE edges"

            # Verify both edges point to the SAME ResourceVersion (only 1 version node)
            result = await session.run(
                """
                MATCH (a:AgentText)-[:READ_FILE]->(v:ResourceVersion)
                RETURN count(DISTINCT v) as version_count
                """
            )
            record = await result.single()
            assert record["version_count"] == 1, "Both edges should point to same ResourceVersion"

            # Verify edges come from different conversations
            result = await session.run(
                """
                MATCH (a:AgentText)-[:READ_FILE]->(v:ResourceVersion)
                RETURN collect(DISTINCT a.conversation_id) as conversations
                """
            )
            record = await result.single()
            assert set(record["conversations"]) == {"conv-1", "conv-2"}

    async def test_conversations_connected_via_shared_resource(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify that a graph path exists between two conversations via shared resource.

        This is the key property: conversations that touch the same resource should
        be connected in the graph, enabling cross-conversation knowledge retrieval.
        """
        shared_file = tmp_path / "config.py"
        content = "DEBUG = True"

        # Two conversations read the same file
        for conv_id in ["conv-alice", "conv-bob"]:
            messages = [
                Message(role="user", content=f"Read config ({conv_id})"),
                Message(
                    role="assistant",
                    content="Reading...",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{conv_id}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{conv_id}"),
            ]
            await tracemem_integration.import_trace(conv_id, messages)

        # Verify a graph path exists between the two conversation's AgentText nodes
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (a1:AgentText {conversation_id: 'conv-alice'})
                      -[:READ_FILE]->(v:ResourceVersion)<-[:READ_FILE]-
                      (a2:AgentText {conversation_id: 'conv-bob'})
                RETURN count(*) as path_exists
                """
            )
            record = await result.single()
            assert record["path_exists"] == 1, (
                "Both conversations should be connected via shared ResourceVersion"
            )

    async def test_version_of_edge_exists_for_shared_resource(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify VERSION_OF edge links ResourceVersion to Resource hypernode."""
        shared_file = tmp_path / "utils.py"
        content = "def helper(): pass"

        # Two conversations read the same file
        for conv_id in ["conv-1", "conv-2"]:
            messages = [
                Message(role="user", content="Read utils"),
                Message(
                    role="assistant",
                    content="Reading...",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{conv_id}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{conv_id}"),
            ]
            await tracemem_integration.import_trace(conv_id, messages)

        # Verify VERSION_OF edge exists
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (v:ResourceVersion)-[:VERSION_OF]->(r:Resource)
                RETURN count(*) as version_of_count,
                       count(DISTINCT r) as resource_count,
                       count(DISTINCT v) as version_count
                """
            )
            record = await result.single()
            assert record["version_of_count"] == 1, "Should have exactly one VERSION_OF edge"
            assert record["resource_count"] == 1, "Should have exactly one Resource"
            assert record["version_count"] == 1, "Should have exactly one ResourceVersion"

    async def test_three_conversations_all_connected_to_same_resource(
        self, tracemem_integration: TraceMem, tmp_path
    ):
        """Verify three or more conversations all create edges to shared resource."""
        shared_file = tmp_path / "main.py"
        content = "if __name__ == '__main__': main()"

        # Three conversations read the same file
        for i in range(3):
            messages = [
                Message(role="user", content=f"Read main.py - request {i}"),
                Message(
                    role="assistant",
                    content="Reading main.py...",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{i}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{i}"),
            ]
            await tracemem_integration.import_trace(f"conv-{i}", messages)

        # Verify all three have edges
        async with tracemem_integration._graph_store._driver.session(
            database="neo4j"
        ) as session:
            result = await session.run(
                """
                MATCH (a:AgentText)-[r:READ_FILE]->(v:ResourceVersion)
                RETURN count(r) as edge_count,
                       collect(DISTINCT a.conversation_id) as conversations
                """
            )
            record = await result.single()
            assert record["edge_count"] == 3
            assert set(record["conversations"]) == {"conv-0", "conv-1", "conv-2"}


class TestRetrievalStrategy:
    """Tests for retrieval strategy and context retrieval."""

    async def test_get_context_returns_complete_chain(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify get_context returns complete chain."""
        test_file = tmp_path / "context_test.py"
        test_file.write_text("def test(): pass")

        messages = [
            Message(role="user", content="What does context_test.py do?"),
            Message(
                role="assistant",
                content="Let me read it for you.",
                tool_calls=[
                    ToolCall(id="c1", name="read_file", args={"path": str(test_file)})
                ],
            ),
            Message(role="tool", content="def test(): pass", tool_call_id="c1"),
        ]
        result = await tracemem_integration.import_trace("conv-1", messages)

        context = await retrieval_strategy.get_context(result["user_text"])

        # Verify complete chain
        assert context.user_text.text == "What does context_test.py do?"
        assert context.agent_text.text == "Let me read it for you."
        assert len(context.tool_uses) == 1

        tool_use = context.tool_uses[0]
        assert tool_use.tool_name == "READ_FILE"
        assert tool_use.resource_version is not None
        assert tool_use.resource_version.content_hash == hashlib.sha256(
            "def test(): pass".encode()
        ).hexdigest()

    async def test_search_finds_relevant_messages(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify search returns relevant results."""
        # Add some messages
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="How do I fix the authentication bug?")
        )
        await tracemem_integration.add_message(
            "conv-2", Message(role="user", content="What's the weather today?")
        )
        await tracemem_integration.add_message(
            "conv-3", Message(role="user", content="Debug the login issue")
        )

        # Search for authentication-related content
        results = await retrieval_strategy.search(
            "authentication problem",
            config=RetrievalConfig(limit=5),
        )

        assert len(results) >= 1
        # Results should be ordered by relevance

    async def test_search_excludes_conversation(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify search excludes specified conversation."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Test message one")
        )
        await tracemem_integration.add_message(
            "conv-2", Message(role="user", content="Test message two")
        )

        # Search excluding conv-1
        results = await retrieval_strategy.search(
            "Test message",
            config=RetrievalConfig(exclude_conversation_id="conv-1", limit=10),
        )

        # Should not include conv-1
        assert all(r.conversation_id != "conv-1" for r in results)

    async def test_search_with_context(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify search can include full context."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Help with code")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Sure, I can help!")
        )

        results = await retrieval_strategy.search(
            "code help",
            config=RetrievalConfig(limit=5, include_context=True),
        )

        assert len(results) >= 1
        # First result should have context populated
        result = results[0]
        assert result.context is not None
        assert result.context.user_text is not None


class TestRetrievalConfig:
    """Tests for RetrievalConfig-based retrieval."""

    async def test_search_with_config(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify search works with RetrievalConfig."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="How do I authenticate users?")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="Use JWT tokens.")
        )

        config = RetrievalConfig(limit=5, include_context=True)
        results = await retrieval_strategy.search("authentication", config=config)

        assert len(results) >= 1
        assert results[0].context is not None

    async def test_search_with_config_excludes_conversation(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify config.exclude_conversation_id works."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Message in conv 1")
        )
        await tracemem_integration.add_message(
            "conv-2", Message(role="user", content="Message in conv 2")
        )

        config = RetrievalConfig(exclude_conversation_id="conv-1")
        results = await retrieval_strategy.search("Message", config=config)

        assert all(r.conversation_id != "conv-1" for r in results)

    async def test_search_with_include_context(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify search with include_context=True populates context."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Code review request")
        )
        await tracemem_integration.add_message(
            "conv-1", Message(role="assistant", content="I'll review it.")
        )

        config = RetrievalConfig(include_context=True)
        results = await retrieval_strategy.search("review", config=config)

        assert len(results) >= 1
        assert results[0].context is not None

    async def test_get_conversations_for_resource_with_limit(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify limit parameter works for resource queries."""
        shared_file = tmp_path / "limited.py"
        content = "# limited test"

        # Create 5 conversations reading the same file
        for i in range(5):
            messages = [
                Message(role="user", content=f"Read limited.py - request {i}"),
                Message(
                    role="assistant",
                    content="Reading...",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{i}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{i}"),
            ]
            await tracemem_integration.import_trace(f"conv-{i}", messages)

        canonical_uri = _canonicalize_file_uri(f"file://{shared_file}", root=None)

        # Request only 3
        config = RetrievalConfig(limit=3)
        conversations = await retrieval_strategy.get_conversations_for_resource(
            canonical_uri, config=config
        )

        assert len(conversations) == 3

    async def test_get_conversations_for_resource_sort_order(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify sort_order parameter affects result order."""
        shared_file = tmp_path / "sorted.py"
        content = "# sorted test"

        # Create 3 conversations
        for i in range(3):
            messages = [
                Message(role="user", content=f"Read sorted.py - request {i}"),
                Message(
                    role="assistant",
                    content="Reading...",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{i}",
                            name="read_file",
                            args={"path": str(shared_file)},
                        )
                    ],
                ),
                Message(role="tool", content=content, tool_call_id=f"call_{i}"),
            ]
            await tracemem_integration.import_trace(f"conv-{i}", messages)

        canonical_uri = _canonicalize_file_uri(f"file://{shared_file}", root=None)

        # Get descending (default - most recent first)
        config_desc = RetrievalConfig(sort_order="desc")
        conversations_desc = await retrieval_strategy.get_conversations_for_resource(
            canonical_uri, config=config_desc
        )

        # Get ascending (oldest first)
        config_asc = RetrievalConfig(sort_order="asc")
        conversations_asc = await retrieval_strategy.get_conversations_for_resource(
            canonical_uri, config=config_asc
        )

        # Order should be reversed
        desc_ids = [c.conversation_id for c in conversations_desc]
        asc_ids = [c.conversation_id for c in conversations_asc]
        assert desc_ids == list(reversed(asc_ids))

    async def test_conversation_reference_includes_agent_text(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify ConversationReference includes agent_text field."""
        shared_file = tmp_path / "agent_response.py"

        messages = [
            Message(role="user", content="Read the agent_response file"),
            Message(
                role="assistant",
                content="Here is the file content analysis.",
                tool_calls=[
                    ToolCall(
                        id="call_1",
                        name="read_file",
                        args={"path": str(shared_file)},
                    )
                ],
            ),
            Message(role="tool", content="# agent response test", tool_call_id="call_1"),
        ]
        await tracemem_integration.import_trace("conv-agent", messages)

        canonical_uri = _canonicalize_file_uri(f"file://{shared_file}", root=None)
        conversations = await retrieval_strategy.get_conversations_for_resource(
            canonical_uri
        )

        assert len(conversations) == 1
        assert conversations[0].agent_text == "Here is the file content analysis."

    async def test_get_conversations_for_resource_with_config(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
        tmp_path,
    ):
        """Verify get_conversations_for_resource works with config."""
        shared_file = tmp_path / "history.py"

        messages = [
            Message(role="user", content="Read history file"),
            Message(
                role="assistant",
                content="Reading history...",
                tool_calls=[
                    ToolCall(
                        id="call_h",
                        name="read_file",
                        args={"path": str(shared_file)},
                    )
                ],
            ),
            Message(role="tool", content="# history test", tool_call_id="call_h"),
        ]
        await tracemem_integration.import_trace("conv-history", messages)

        canonical_uri = _canonicalize_file_uri(f"file://{shared_file}", root=None)

        config = RetrievalConfig(limit=5)
        conversations = await retrieval_strategy.get_conversations_for_resource(
            canonical_uri, config=config
        )

        assert len(conversations) == 1
        assert conversations[0].conversation_id == "conv-history"

    async def test_search_with_exclude_conversation_in_config(
        self,
        tracemem_integration: TraceMem,
        retrieval_strategy: HybridRetrievalStrategy,
    ):
        """Verify exclude_conversation_id in config filters results."""
        await tracemem_integration.add_message(
            "conv-1", Message(role="user", content="Config exclude test message 1")
        )
        await tracemem_integration.add_message(
            "conv-2", Message(role="user", content="Config exclude test message 2")
        )

        config = RetrievalConfig(limit=10, exclude_conversation_id="conv-1")
        results = await retrieval_strategy.search(
            "Config exclude test",
            config=config,
        )

        assert all(r.conversation_id != "conv-1" for r in results)
