"""Tests for graph-based retrieval (MemGraphRAG pattern).

Tests entity relationship graph traversal and document retrieval
through co-occurrence and named entity relationships.
"""

import os
from pathlib import Path
import time

import pytest

from cyberbrain26.data_plane.sqlite_store import SQLiteDatabase
from cyberbrain26.knowledge_plane.contracts import (
    KnowledgeDocument,
    DocumentStatus,
)
from cyberbrain26.knowledge_plane.memory_service import MemoryService
from cyberbrain26.knowledge_plane.navigation import KnowledgeNavigator


@pytest.fixture
def ms() -> MemoryService:
    """Provide a MemoryService with isolated env."""
    from cyberbrain26.knowledge_plane.embedding_config import EmbeddingConfig
    from cyberbrain26.security_plane.audit import AuditLog

    data_dir = Path(os.environ.get("CYBERBRAIN_DATA_DIR", "/tmp"))
    db_path = Path(os.environ.get("CYBERBRAIN_DB_PATH", "/tmp/cyberbrain.sqlite3"))

    config = EmbeddingConfig(vector_dim=1024, lancedb_directory=data_dir / "lancedb")
    audit_log = AuditLog(database=SQLiteDatabase(str(db_path)))

    return MemoryService(
        database=SQLiteDatabase(db_path),
        audit_log=audit_log,
        config=config,
        auto_classify=False,
    )


def test_memory_service_graph_search_basic(ms: MemoryService):
    """Test that MemoryService.graph_search finds documents via entity traversal."""
    doc1, _ = ms.register_text(
        content=b"Apple develops iPhone technology",
        source_uri="memory://doc1",
        source_id="test-entities",
        topic="technology",
        tags=["apple", "iphone"],
        entities=["apple", "iphone"],
    )

    doc2, _ = ms.register_text(
        content=b"Microsoft releases Windows OS",
        source_uri="memory://doc2",
        source_id="test-entities",
        topic="technology",
        tags=["microsoft", "windows"],
        entities=["microsoft", "windows"],
    )

    doc3, _ = ms.register_text(
        content=b"Google search engine AI",
        source_uri="memory://doc3",
        source_id="test-entities",
        topic="ai",
        tags=["google", "search"],
        entities=["google", "search"],
    )

    results = ms.graph_search("Apple", limit=10, depth=2)
    assert len(results) >= 1

    for r in results:
        assert "id" in r
        assert "score" in r


def test_memory_service_graph_search_no_match(ms: MemoryService):
    """Test that graph_search returns empty when no matching entities."""
    results = ms.graph_search("NonExistentEntity123", limit=10, depth=2)
    assert len(results) == 0


def test_memory_service_graph_search_depth_control(ms: MemoryService):
    """Test that graph_search depth parameter controls traversal breadth."""
    doc1, _ = ms.register_text(
        content=b"Apple Microsoft partnership AI research",
        source_uri="memory://doc1",
        source_id="test-entities",
        topic="technology",
        tags=["apple", "microsoft", "ai"],
        entities=["apple", "microsoft", "ai", "partnership"],
    )

    results_shallow = ms.graph_search("Apple", limit=10, depth=1)
    results_deep = ms.graph_search("Apple", limit=10, depth=2)

    assert len(results_shallow) >= 1
    assert len(results_deep) >= 1


def test_knowledge_navigator_entity_graph():
    """Test that KnowledgeNavigator correctly maintains entity graph."""
    navigator = KnowledgeNavigator()

    doc1 = KnowledgeDocument(
        document_id="doc1",
        source_id="test",
        source_uri="file:///doc1.txt",
        content_sha256="abc",
        media_type="text/plain",
        trust="reviewed",
        sensitivity="personal",
        status=DocumentStatus.INDEXED,
        metadata={
            "topic": "technology",
            "tags": "apple,microsoft",
            "entities": "apple,microsoft,partnership",
        },
        created_at=time.time(),
        updated_at=time.time(),
        valid_from=time.time(),
        valid_to=0.0,
        superseded_by="",
    )

    doc2 = KnowledgeDocument(
        document_id="doc2",
        source_id="test",
        source_uri="file:///doc2.txt",
        content_sha256="def",
        media_type="text/plain",
        trust="reviewed",
        sensitivity="personal",
        status=DocumentStatus.INDEXED,
        metadata={
            "topic": "technology",
            "tags": "microsoft,windows",
            "entities": "microsoft,windows,partnership",
        },
        created_at=time.time(),
        updated_at=time.time(),
        valid_from=time.time(),
        valid_to=0.0,
        superseded_by="",
    )

    navigator.add_document(doc1)
    navigator.add_document(doc2)

    assert "apple" in navigator.get_all_tags()
    assert "microsoft" in navigator.get_all_tags()
    assert "partnership" in navigator.get_all_tags()

    apple_neighbors = navigator.get_entity_neighbors("apple")
    assert "microsoft" in apple_neighbors
    assert "partnership" in apple_neighbors

    # get_entity_documents returns direct associations only
    apple_docs = navigator.get_entity_documents("apple")
    assert "doc1" in apple_docs
    assert "doc2" not in apple_docs  # doc2 doesn't have "apple" entity

    # microsoft appears in both docs
    microsoft_docs = navigator.get_entity_documents("microsoft")
    assert "doc1" in microsoft_docs
    assert "doc2" in microsoft_docs

    # partnership appears in both docs
    partnership_docs = navigator.get_entity_documents("partnership")
    assert "doc1" in partnership_docs
    assert "doc2" in partnership_docs


def test_memory_service_graph_search_with_relationship_type(ms: MemoryService):
    """Test graph_search with optional relationship_type filter."""
    doc1, _ = ms.register_text(
        content=b"Apple and Microsoft partnership 2024",
        source_uri="memory://doc1",
        source_id="test-entities",
        topic="technology",
        tags=["apple", "microsoft", "partnership"],
        entities=["apple", "microsoft", "partnership"],
    )

    results = ms.graph_search("Apple", limit=10, depth=2, relationship_type="co_reference")
    assert len(results) >= 1