"""Unit tests for the lightweight RAG/KAG demonstration module."""

from services.knowledge_pipeline import (
    KnowledgeGraph,
    KnowledgeWorkflowOrchestrator,
    RAGPipeline,
    run_demo,
)


def test_rag_pipeline_indexes_and_retrieves_documents() -> None:
    pipeline = RAGPipeline(chunk_token_limit=20)
    pipeline.index({
        "doc-1": "RAG pipelines rely on document indexing and vector stores.",
        "doc-2": "Knowledge graphs organise entities for reasoning.",
    })

    results = pipeline.search("How do RAG pipelines work?", limit=1)

    assert len(results) == 1
    assert results[0].doc_id == "doc-1"
    assert "vector stores" in results[0].text.lower()


def test_knowledge_graph_returns_related_edges() -> None:
    graph = KnowledgeGraph()
    graph.add_fact("rag", "stores_in", "vector_db", evidence="Vector DB keeps embeddings queryable.")
    graph.add_fact("kag", "aligns", "knowledge_graph", evidence="Knowledge graph maintains governance.")

    matches = graph.reason("Which vector_db backs the rag pipeline?")

    assert len(matches) == 1
    assert matches[0].source == "rag"
    assert matches[0].target == "vector_db"


def test_orchestrator_combines_rag_and_graph_results() -> None:
    pipeline = RAGPipeline(chunk_token_limit=40)
    pipeline.index({
        "doc": "Cache augmented generation preloads knowledge so retrieval stays instant.",
    })
    graph = KnowledgeGraph()
    graph.add_fact("agent_loop", "includes", "reflection", evidence="Reflection improves outcomes.")

    orchestrator = KnowledgeWorkflowOrchestrator(rag=pipeline, graph=graph)
    answer = orchestrator.answer("How does the agent loop include reflection for cache workflows?", limit=1)

    assert answer.rag_context
    assert "cache augmented" in answer.rag_context[0].lower()
    assert any(edge.target == "reflection" for edge in answer.knowledge_edges)


def test_run_demo_uses_repository_datasets() -> None:
    answer = run_demo("Which workflows activate the moe router?", limit=2)

    assert answer.rag_context  # dataset-backed narrative should surface
    # Demo graph encodes a moe_router -> stripe edge which should be returned
    assert any(edge.source == "moe_router" for edge in answer.knowledge_edges)